diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 651ce856..9159844d 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -542,6 +542,7 @@ Audiobooks only. AudiobookBay lists it for most releases. In Prowlarr results it | `AUTO_OPEN_DOWNLOADS_SIDEBAR` | Automatically open the downloads sidebar when a new download is queued. | boolean | `false` | | `DOWNLOAD_TO_BROWSER_CONTENT_TYPES` | Automatically download completed files to your browser for the selected content types. | string (comma-separated) | _empty list_ | | `MAX_CONCURRENT_DOWNLOADS` | Maximum number of simultaneous downloads. | number | `3` | +| `MAX_CONCURRENT_DOWNLOADS_PER_SOURCE` | Simultaneous downloads allowed from each source (Soulseek, AudiobookBay, Direct Download and so on). A source at its limit does not hold the others back, so a Soulseek download waiting in a peer's queue cannot stop a torrent from starting. 0 uses the single limit above for all sources together. | number | `0` | | `STATUS_TIMEOUT` | How long to keep completed/failed downloads in the queue display. | number | `3600` |
@@ -852,6 +853,17 @@ Maximum number of simultaneous downloads. - **Requires restart:** Yes - **Constraints:** min: 1, max: 10 +#### `MAX_CONCURRENT_DOWNLOADS_PER_SOURCE` + +**Max Concurrent Downloads Per Source** + +Simultaneous downloads allowed from each source (Soulseek, AudiobookBay, Direct Download and so on). A source at its limit does not hold the others back, so a Soulseek download waiting in a peer's queue cannot stop a torrent from starting. 0 uses the single limit above for all sources together. + +- **Type:** number +- **Default:** `0` +- **Requires restart:** Yes +- **Constraints:** min: 0, max: 10 + #### `STATUS_TIMEOUT` **Status Timeout (seconds)** diff --git a/shelfmark/config/settings.py b/shelfmark/config/settings.py index ed6b559c..f26086d1 100644 --- a/shelfmark/config/settings.py +++ b/shelfmark/config/settings.py @@ -1494,6 +1494,20 @@ def download_settings() -> list[SettingsField]: max_value=10, requires_restart=True, ), + NumberField( + key="MAX_CONCURRENT_DOWNLOADS_PER_SOURCE", + label="Max Concurrent Downloads Per Source", + description=( + "Simultaneous downloads allowed from each source (Soulseek, AudiobookBay, " + "Direct Download and so on). A source at its limit does not hold the others " + "back, so a Soulseek download waiting in a peer's queue cannot stop a torrent " + "from starting. 0 uses the single limit above for all sources together." + ), + default=0, + min_value=0, + max_value=10, + requires_restart=True, + ), NumberField( key="STATUS_TIMEOUT", label="Status Timeout (seconds)", diff --git a/shelfmark/core/models.py b/shelfmark/core/models.py index ff7ae9cc..91657197 100644 --- a/shelfmark/core/models.py +++ b/shelfmark/core/models.py @@ -163,6 +163,10 @@ class DownloadTask: download_paths: list[str] = field(default_factory=list) last_error_message: str | None = None last_error_type: str | None = None + + # True while the download client has this task's torrent in its own queue, waiting for a + # free slot. Such a task does no work yet, so it does not count against its source's limit. + client_queued: bool = False staged_path: str | None = None def __lt__(self, other: DownloadTask) -> bool: diff --git a/shelfmark/core/queue.py b/shelfmark/core/queue.py index 8567c8d6..99a69b75 100644 --- a/shelfmark/core/queue.py +++ b/shelfmark/core/queue.py @@ -92,12 +92,23 @@ class BookQueue: logger.warning("Queue hook failed while adding task %s: %s", task_id, exc) return True - def get_next(self) -> tuple[str, Event] | None: - """Get next task ID from queue with cancellation flag.""" - # Use iterative approach to avoid stack overflow if many items are cancelled - while True: - try: - queue_item = self._queue.get_nowait() + def get_next( + self, accept: Callable[[DownloadTask], bool] | None = None + ) -> tuple[str, Event] | None: + """Get next task ID from queue with cancellation flag. + + ``accept`` lets the caller pass over tasks it cannot start yet (their source is at its + limit). A task passed over keeps its place: it goes back with the same priority and + arrival time, so it is first in line once there is room. + """ + passed_over: list[QueueItem] = [] + try: + # Use iterative approach to avoid stack overflow if many items are cancelled + while True: + try: + queue_item = self._queue.get_nowait() + except queue.Empty: + return None task_id = queue_item.book_id # QueueItem uses book_id as the ID field with self._lock: @@ -105,14 +116,19 @@ class BookQueue: if task_id in self._status and self._status[task_id] == QueueStatus.CANCELLED: continue # Skip cancelled items, try next + task = self._task_data.get(task_id) + if accept is not None and task is not None and not accept(task): + passed_over.append(queue_item) + continue + # Create cancellation flag for this download cancel_flag = Event() self._cancel_flags[task_id] = cancel_flag self._active_downloads[task_id] = True - except queue.Empty: - return None - else: return task_id, cancel_flag + finally: + for item in passed_over: + self._queue.put(item) def get_task(self, task_id: str) -> DownloadTask | None: """Get a task by its ID.""" diff --git a/shelfmark/download/clients/base_handler.py b/shelfmark/download/clients/base_handler.py index b8275db2..26c91b1d 100644 --- a/shelfmark/download/clients/base_handler.py +++ b/shelfmark/download/clients/base_handler.py @@ -1095,6 +1095,7 @@ class ExternalClientHandler(DownloadHandler, ABC): while not cancel_flag.is_set(): status = client.get_status(download_id) progress_callback(status.progress) + task.client_queued = status.state == DownloadState.QUEUED # Check for completion if status.complete: diff --git a/shelfmark/download/orchestrator.py b/shelfmark/download/orchestrator.py index 31b74669..b29bc5ed 100644 --- a/shelfmark/download/orchestrator.py +++ b/shelfmark/download/orchestrator.py @@ -10,6 +10,7 @@ import threading import time from concurrent.futures import Future, ThreadPoolExecutor from email.utils import parseaddr +from functools import partial from pathlib import Path from threading import Event, Lock from typing import TYPE_CHECKING, Any @@ -1168,13 +1169,45 @@ def _cancel_stalled_task(task_id: str) -> None: book_queue.cancel_download(task_id) +# With a per-source limit the pool must hold every source's downloads at once. +_MAX_SOURCES = 8 + + +def _per_source_limit() -> int: + """Downloads allowed at once from one source, or 0 when only the global limit applies.""" + return normalize_positive_int(config.get("MAX_CONCURRENT_DOWNLOADS_PER_SOURCE", 0)) or 0 + + +def _source_has_room(task: DownloadTask, *, running: dict[str | None, int], limit: int) -> bool: + """Whether the task's source is below its limit of simultaneous downloads.""" + return running.get(task.source, 0) < limit + + +def _working_source_of(task_id: str) -> str | None: + """The source a running task counts against, or None when it is only waiting in a client queue.""" + task = book_queue.get_task(task_id) + if task is None or task.client_queued: + return None + return task.source + + +def _source_of(task_id: str) -> str | None: + task = book_queue.get_task(task_id) + return task.source if task is not None else None + + def concurrent_download_loop() -> None: """Run the main concurrent download coordinator.""" max_workers = normalize_positive_int(config.MAX_CONCURRENT_DOWNLOADS) or 1 + per_source = _per_source_limit() + pool_size = per_source * _MAX_SOURCES if per_source else max_workers main_loop_sleep_time = _config_float(config.MAIN_LOOP_SLEEP_TIME, 0.5) - logger.info("Starting concurrent download loop with %s workers", max_workers) + if per_source: + logger.info("Starting concurrent download loop with %s downloads per source", per_source) + else: + logger.info("Starting concurrent download loop with %s workers", max_workers) - with ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="Download") as executor: + with ThreadPoolExecutor(max_workers=pool_size, thread_name_prefix="Download") as executor: active_futures: dict[Future, tuple[str, Event]] = {} # Track active download futures stalled_tasks: set[str] = set() # Track tasks already cancelled due to stall consecutive_errors = 0 @@ -1238,20 +1271,37 @@ def concurrent_download_loop() -> None: stalled_tasks.add(task_id) # Start new downloads if we have capacity - while len(active_futures) < max_workers: - next_download = book_queue.get_next() + while len(active_futures) < pool_size: + running_by_source: dict[str | None, int] = {} + if per_source: + for running_id, _flag in active_futures.values(): + key = _working_source_of(running_id) + if key is not None: + running_by_source[key] = running_by_source.get(key, 0) + 1 + accept = ( + partial(_source_has_room, running=running_by_source, limit=per_source) + if per_source + else None + ) + next_download = book_queue.get_next(accept) if accept else book_queue.get_next() if not next_download: break - # Stagger concurrent downloads to avoid rate limiting on shared download servers - # Only delay if other downloads are already active - if active_futures: + task_id, cancel_flag = next_download + + # Stagger concurrent downloads to avoid rate limiting on shared download + # servers. With per-source limits only a download from the same source + # shares a server, so only those wait for each other. + sharing_a_server = ( + running_by_source.get(_source_of(task_id), 0) > 0 + if per_source + else bool(active_futures) + ) + if sharing_a_server: stagger_delay = _RNG.uniform(2, 5) logger.debug("Staggering download start by %.1fs", stagger_delay) time.sleep(stagger_delay) - task_id, cancel_flag = next_download - # Submit download job to thread pool future = executor.submit(_process_single_download, task_id, cancel_flag) active_futures[future] = (task_id, cancel_flag) diff --git a/tests/core/test_queue_passing_over.py b/tests/core/test_queue_passing_over.py new file mode 100644 index 00000000..181e8909 --- /dev/null +++ b/tests/core/test_queue_passing_over.py @@ -0,0 +1,52 @@ +"""A queue read that passes over the tasks its caller cannot start yet.""" + +from shelfmark.core.models import DownloadTask, QueueStatus +from shelfmark.core.queue import BookQueue + + +def _task_from(source: str, task_id: str, priority: int = 0) -> DownloadTask: + return DownloadTask(task_id=task_id, source=source, title=task_id, priority=priority) + + +def test_get_next_can_pass_over_a_task_and_leave_it_first_in_line(): + queue = BookQueue() + assert queue.add(_task_from("slskd", "first-soulseek")) is True + assert queue.add(_task_from("audiobookbay", "torrent")) is True + assert queue.add(_task_from("slskd", "second-soulseek")) is True + + picked = queue.get_next(lambda task: task.source != "slskd") + + assert picked is not None + assert picked[0] == "torrent" + # Both Soulseek tasks are back, in their original order. + assert queue.get_next()[0] == "first-soulseek" + assert queue.get_next()[0] == "second-soulseek" + assert queue.get_next() is None + + +def test_get_next_returns_nothing_when_every_task_is_passed_over_and_keeps_them(): + queue = BookQueue() + assert queue.add(_task_from("slskd", "a")) is True + assert queue.add(_task_from("slskd", "b")) is True + + assert queue.get_next(lambda _task: False) is None + + assert queue.get_next()[0] == "a" + assert queue.get_next()[0] == "b" + + +def test_a_passed_over_task_keeps_its_priority_against_a_later_one(): + queue = BookQueue() + assert queue.add(_task_from("slskd", "urgent", priority=-10)) is True + assert queue.add(_task_from("audiobookbay", "normal", priority=0)) is True + + assert queue.get_next(lambda task: task.source != "slskd")[0] == "normal" + assert queue.get_next()[0] == "urgent" + + +def test_a_cancelled_task_is_dropped_even_when_a_filter_is_given(): + queue = BookQueue() + assert queue.add(_task_from("audiobookbay", "gone")) is True + queue.update_status("gone", QueueStatus.CANCELLED) + + assert queue.get_next(lambda _task: True) is None diff --git a/tests/download/test_client_queued.py b/tests/download/test_client_queued.py new file mode 100644 index 00000000..7f603ac5 --- /dev/null +++ b/tests/download/test_client_queued.py @@ -0,0 +1,51 @@ +"""A torrent queued inside its client is not working yet, and the task says so.""" + +from __future__ import annotations + +from threading import Event +from unittest.mock import MagicMock, patch + +from shelfmark.core.models import DownloadTask +from shelfmark.download.clients import DownloadState, DownloadStatus +from shelfmark.release_sources.prowlarr.handler import ProwlarrHandler + + +def _status(state: DownloadState) -> DownloadStatus: + return DownloadStatus(progress=0.0, state=state, message=None, complete=False, file_path=None) + + +def _poll(states: list[DownloadState]) -> tuple[DownloadTask, list[bool]]: + cancel = Event() + seen = iter(states) + flags: list[bool] = [] + task = DownloadTask(task_id="t", source="prowlarr", title="Book") + client = MagicMock() + client.name = "qbittorrent" + + def get_status(_id: str) -> DownloadStatus: + flags.append(task.client_queued) # what the previous poll left behind + try: + return _status(next(seen)) + except StopIteration: + cancel.set() + return _status(states[-1]) + + client.get_status.side_effect = get_status + with patch.object(ProwlarrHandler, "_poll_interval", return_value=0.0): + ProwlarrHandler()._poll_and_complete( + client, "hash", "torrent", task, cancel, lambda _p: None, lambda *_a: None + ) + return task, flags + + +def test_the_task_is_marked_while_its_torrent_waits_in_the_clients_queue(): + queued, downloading = DownloadState.QUEUED, DownloadState.DOWNLOADING + _task, flags = _poll([queued, queued, downloading, downloading]) + + assert flags[:5] == [False, True, True, False, False] + + +def test_a_task_that_never_waited_is_never_marked(): + _task, flags = _poll([DownloadState.DOWNLOADING] * 3) + + assert not any(flags) diff --git a/tests/download/test_orchestrator_lifecycle.py b/tests/download/test_orchestrator_lifecycle.py index 02e825f5..f10e87a3 100644 --- a/tests/download/test_orchestrator_lifecycle.py +++ b/tests/download/test_orchestrator_lifecycle.py @@ -305,3 +305,120 @@ def test_download_task_completes_blackhole_handoff_without_post_processing( call(task.task_id, "Cancelled, but the torrent was already handed off to blackhole") in queue.update_status_message.call_args_list ) + + +class _PoolRecorder: + """Stands in for the thread pool: remembers what was started and never finishes it.""" + + def __init__(self, *args, **kwargs) -> None: + self.max_workers = kwargs.get("max_workers") + _PoolRecorder.last = self + self.started: list[str] = [] + + def __enter__(self) -> _PoolRecorder: + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + return False + + def submit(self, fn, task_id, cancel_flag): + from concurrent.futures import Future + + self.started.append(task_id) + return Future() # never done, so the download stays running + + +def _run_one_pass(monkeypatch, tasks, *, per_source, global_limit=3): + import shelfmark.download.orchestrator as orchestrator + from shelfmark.core.queue import BookQueue + + queue = BookQueue() + for task in tasks: + assert queue.add(task) is True + values = { + "MAX_CONCURRENT_DOWNLOADS_PER_SOURCE": per_source, + "MAX_CONCURRENT_DOWNLOADS": global_limit, + "MAIN_LOOP_SLEEP_TIME": 0.0, + } + monkeypatch.setattr(orchestrator, "book_queue", queue) + monkeypatch.setattr(orchestrator, "ThreadPoolExecutor", _PoolRecorder) + + def sleep(delay): + if delay == 0.0: # the loop's own pause ends the pass; a stagger delay is just skipped + raise _StopLoop + + monkeypatch.setattr(orchestrator.time, "sleep", sleep) + monkeypatch.setattr(orchestrator.config, "get", lambda k, d=None: values.get(k, d)) + monkeypatch.setattr( + orchestrator.config, "MAX_CONCURRENT_DOWNLOADS", global_limit, raising=False + ) + monkeypatch.setattr(orchestrator.config, "MAIN_LOOP_SLEEP_TIME", 0.0, raising=False) + with pytest.raises(_StopLoop): + orchestrator.concurrent_download_loop() + return _PoolRecorder.last + + +def _t(source, task_id): + return DownloadTask(task_id=task_id, source=source, title=task_id) + + +def test_a_source_at_its_limit_does_not_hold_back_the_others(monkeypatch): + pool = _run_one_pass( + monkeypatch, + [_t("slskd", "s1"), _t("slskd", "s2"), _t("slskd", "s3"), _t("audiobookbay", "torrent")], + per_source=1, + ) + + assert sorted(pool.started) == ["s1", "torrent"] + + +def test_each_source_gets_its_own_limit(monkeypatch): + pool = _run_one_pass( + monkeypatch, + [_t("slskd", "s1"), _t("slskd", "s2"), _t("audiobookbay", "a1"), _t("audiobookbay", "a2")], + per_source=2, + ) + + assert sorted(pool.started) == ["a1", "a2", "s1", "s2"] + + +def test_with_no_per_source_limit_the_one_global_limit_applies(monkeypatch): + pool = _run_one_pass( + monkeypatch, + [_t("slskd", "s1"), _t("slskd", "s2"), _t("audiobookbay", "a1"), _t("audiobookbay", "a2")], + per_source=0, + global_limit=3, + ) + + assert len(pool.started) == 3 + + +def test_the_pool_is_big_enough_for_every_source_at_its_limit(monkeypatch): + import shelfmark.download.orchestrator as orchestrator + + pool = _run_one_pass(monkeypatch, [], per_source=2) + + assert pool.max_workers == 2 * orchestrator._MAX_SOURCES + + +def test_a_source_has_room_only_below_its_limit(): + import shelfmark.download.orchestrator as orchestrator + + task = _t("slskd", "x") + + assert orchestrator._source_has_room(task, running={"slskd": 1}, limit=2) is True + assert orchestrator._source_has_room(task, running={"slskd": 2}, limit=2) is False + assert orchestrator._source_has_room(task, running={"audiobookbay": 5}, limit=2) is True + + +def test_a_task_waiting_in_its_clients_queue_does_not_use_up_its_source(monkeypatch): + waiting = DownloadTask(task_id="waiting", source="audiobookbay", title="w", client_queued=True) + pool = _run_one_pass( + monkeypatch, + [waiting, _t("audiobookbay", "next"), _t("audiobookbay", "third")], + per_source=1, + ) + + # "waiting" is first and takes the one slot; being queued in qBittorrent it gives it back, + # so "next" starts; "next" is working, so "third" has to wait. + assert pool.started == ["waiting", "next"]