mirror of
https://github.com/calibrain/shelfmark.git
synced 2026-11-15 10:28:33 +00:00
feat(download): a concurrency limit per source (#1495)
With one global limit, a few tasks that are only waiting hold every slot. A Soulseek download sitting in a peer's queue is the example I hit: three of those and a torrent that could start waits behind them. `MAX_CONCURRENT_DOWNLOADS_PER_SOURCE` (0 keeps the single global limit, which is the default). With it set, the coordinator starts a download only when its source is below the limit and passes over tasks of a full source without losing their place in the queue. `get_next` gets an optional filter for this, and a cancelled task is still dropped when one is given. The start-up stagger now applies only between downloads from the same source, and the pool is sized so every source can reach its limit. The second commit: qBittorrent keeps its own queue of torrents waiting for a free slot. The task polling one of those does no work, so it no longer counts against its source's limit and the next task of that source can start. The client does the pacing, as it does for its own queue. Tests cover the queue pass-over (order and priority kept, nothing eligible, cancelled tasks), the coordinator (a full source doesn't block the others, each source has its own limit, the global limit when unset) and the client-queue case. Ruff and the backend suite pass.
This commit is contained in:
1 parent
e8297c94a9
commit
06a67a7150
9 files changed
+335
-18
No files matched your search
@@ -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` |
|
||||
|
||||
<details>
|
||||
@@ -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)**
|
||||
|
||||
@@ -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)",
|
||||
|
||||
@@ -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:
|
||||
|
||||
+25
-9
@@ -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."""
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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"]
|
||||
Reference in new issue
Block a user