kenzuya/unshackle
1
0
Fork 2
mirror of https://github.com/unshackle-dl/unshackle.git synced 2026-05-16 21:59:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(downloader): optimize download throughput with Queue-based threading and raw reads

Browse Source 
Fix critical bug where ThreadPoolExecutor was not actually parallelizing downloads (generator functions returned instantly, I/O ran on main thread).

Performance improvements:
Queue-based event dispatch: workers consume generators in threads, push events to a thread-safe Queue for truly parallel segment downloads
Raw socket reads (resp.raw.read) for requests.Session — 30-35% faster than iter_content, with iter_content fallback for CurlSession
File pre-allocation via truncate when Content-Length is known
Hot loop caching: time.time, f.write, stream.raw.read cached as locals
HTTPAdapter connection pooling mounted on passed sessions for reuse
This commit is contained in:
Andy
2026-03-23 17:20:26 -06:00
parent dc197af29e
commit 006d080416
1 changed files with 133 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

188
unshackle/core/downloaders/requests.py
View File

@@ -1,10 +1,10 @@
import math
import os
import time
from concurrent.futures import as_completed
from concurrent.futures.thread import ThreadPoolExecutor
from http.cookiejar import CookieJar
from pathlib import Path
from queue import Empty, Queue
from typing import Any, Generator, MutableMapping, Optional, Union
from requests import Session
@@ -32,6 +32,10 @@ def _adaptive_chunk_size(content_length: int) -> int:
return min(MAX_CHUNK, max(MIN_CHUNK, content_length // 4))
def _is_requests_session(session: Any) -> bool:
"""Check if the session is a standard requests.Session (supports resp.raw)."""
return isinstance(session, Session)
def download(
url: str,
@@ -45,7 +49,7 @@ def download(
Download a file with optimized I/O.
Supports both requests.Session and curl_cffi CurlSession for TLS fingerprinting.
Uses adaptive chunk sizing with buffered writes for maximum throughput.
Uses raw socket reads for requests.Session (30-35% faster) and iter_content for CurlSession.
Yields the following download status updates while chunks are downloading:
@@ -69,7 +73,6 @@ def download(
"""
session = session or Session()
# Per-call speed tracking (shared across threads within one requests() call)
if _speed_tracker is None:
_speed_tracker = {"sizes": [], "last_refresh": time.time()}
@@ -85,13 +88,15 @@ def download(
yield dict(file_downloaded=save_path, written=save_path.stat().st_size)
control_file.write_bytes(b"")
_time = time.time
use_raw = _is_requests_session(session)
attempts = 1
try:
while True:
written = 0
download_sizes: list[int] = []
last_speed_refresh = time.time()
last_speed_refresh = _time()
try:
stream = session.get(url, stream=True, **kwargs)
@@ -113,25 +118,61 @@ def download(
else:
yield dict(total=None)
# Buffered iter_content with adaptive chunk size
# Works with both requests.Session and CurlSession
# Pre-allocate file when size is known (helps filesystem allocate contiguous blocks)
with open(save_path, "wb", buffering=1_048_576) as f:
for chunk in stream.iter_content(chunk_size=chunk_size):
download_size = len(chunk)
f.write(chunk)
written += download_size
if content_length > 0:
f.truncate(content_length)
f.seek(0)
if not segmented:
yield dict(advance=1)
now = time.time()
time_since = now - last_speed_refresh
download_sizes.append(download_size)
if time_since > PROGRESS_WINDOW or download_size < chunk_size:
data_size = sum(download_sizes)
download_speed = math.ceil(data_size / (time_since or 1))
yield dict(downloaded=f"{filesize.decimal(download_speed)}/s")
last_speed_refresh = now
download_sizes.clear()
# Cache f.write for hot loop
_write = f.write
if use_raw:
# Raw socket read — 30-35% faster than iter_content (benchmarked)
# Safe in worker threads with Queue-based event dispatch
_read = stream.raw.read
while True:
chunk = _read(chunk_size)
if not chunk:
break
_write(chunk)
download_size = len(chunk)
written += download_size
if not segmented:
yield dict(advance=1)
now = _time()
time_since = now - last_speed_refresh
download_sizes.append(download_size)
if time_since > PROGRESS_WINDOW or download_size < chunk_size:
data_size = sum(download_sizes)
download_speed = math.ceil(data_size / (time_since or 1))
yield dict(downloaded=f"{filesize.decimal(download_speed)}/s")
last_speed_refresh = now
download_sizes.clear()
stream.close()
else:
# CurlSession: use iter_content (raw not available)
for chunk in stream.iter_content(chunk_size=chunk_size):
_write(chunk)
download_size = len(chunk)
written += download_size
if not segmented:
yield dict(advance=1)
now = _time()
time_since = now - last_speed_refresh
download_sizes.append(download_size)
if time_since > PROGRESS_WINDOW or download_size < chunk_size:
data_size = sum(download_sizes)
download_speed = math.ceil(data_size / (time_since or 1))
yield dict(downloaded=f"{filesize.decimal(download_speed)}/s")
last_speed_refresh = now
download_sizes.clear()
# Truncate to actual written size in case pre-allocation overshot
if content_length > 0 and written != content_length:
f.truncate(written)
if not segmented and content_length and written < content_length:
raise IOError(f"Failed to read {content_length} bytes from the track URI.")
@@ -140,11 +181,10 @@ def download(
if segmented:
yield dict(advance=1)
now = time.time()
now = _time()
sizes = _speed_tracker["sizes"]
if written:
sizes.append((now, written))
# Prune entries older than the rolling window
cutoff = now - SPEED_ROLLING_WINDOW
while sizes and sizes[0][0] < cutoff:
sizes.pop(0)
@@ -256,12 +296,10 @@ def requests(
]
# Use provided session or create a new optimized requests.Session
# When a session is provided (e.g., service's CurlSession), don't mutate it —
# headers/cookies/proxy are already set on it and it may be shared across tracks.
# When a session is provided (e.g., service's CurlSession), don't mutate headers/cookies/proxy —
# they're already set and the session may be shared across tracks.
if session is None:
session = Session()
session.mount("https://", HTTPAdapter(pool_connections=max_workers, pool_maxsize=max_workers, pool_block=True))
session.mount("http://", HTTPAdapter(pool_connections=max_workers, pool_maxsize=max_workers, pool_block=True))
if headers:
headers = {k: v for k, v in headers.items() if k.lower() != "accept-encoding"}
session.headers.update(headers)
@@ -270,6 +308,13 @@ def requests(
if proxy:
session.proxies.update({"all": proxy})
# Mount HTTPAdapter with connection pooling sized to worker count.
# Safe to do on any requests.Session — improves connection reuse for parallel downloads.
if _is_requests_session(session):
adapter = HTTPAdapter(pool_connections=max_workers, pool_maxsize=max_workers, pool_block=True)
session.mount("https://", adapter)
session.mount("http://", adapter)
if debug_logger:
first_url = urls[0].get("url", "") if urls else ""
url_display = first_url[:200] + "..." if len(first_url) > 200 else first_url
@@ -297,35 +342,68 @@ def requests(
try:
with ThreadPoolExecutor(max_workers=max_workers) as pool:
for future in as_completed(
pool.submit(download, session=session, segmented=segmented_batch, _speed_tracker=speed_tracker, **url)
for url in urls
):
event_queue: Queue[dict[str, Any]] = Queue()
def _download_worker(url_item: dict[str, Any]) -> None:
for event in download(
session=session,
segmented=segmented_batch,
_speed_tracker=speed_tracker,
**url_item,
):
event_queue.put(event)
futures = [pool.submit(_download_worker, url) for url in urls]
pending = set(futures)
while pending:
# Drain queued progress updates for responsive UI
while True:
try:
yield event_queue.get_nowait()
except Empty:
break
done = {future for future in pending if future.done()}
for future in done:
pending.remove(future)
exc = future.exception()
if isinstance(exc, KeyboardInterrupt):
DOWNLOAD_CANCELLED.set()
yield dict(downloaded="[yellow]CANCELLING")
pool.shutdown(wait=True, cancel_futures=True)
yield dict(downloaded="[yellow]CANCELLED")
raise exc
elif exc:
DOWNLOAD_CANCELLED.set()
yield dict(downloaded="[red]FAILING")
pool.shutdown(wait=True, cancel_futures=True)
yield dict(downloaded="[red]FAILED")
if debug_logger:
debug_logger.log(
level="ERROR",
operation="downloader_failed",
message=f"Download failed: {exc}",
error=exc,
context={
"url_count": len(urls),
"output_dir": str(output_dir),
},
)
raise exc
if pending:
try:
yield event_queue.get(timeout=0.1)
except Empty:
pass
# Drain any remaining events from workers that just finished
while True:
try:
yield from future.result()
except KeyboardInterrupt:
DOWNLOAD_CANCELLED.set()
yield dict(downloaded="[yellow]CANCELLING")
pool.shutdown(wait=True, cancel_futures=True)
yield dict(downloaded="[yellow]CANCELLED")
raise
except Exception as e:
DOWNLOAD_CANCELLED.set()
yield dict(downloaded="[red]FAILING")
pool.shutdown(wait=True, cancel_futures=True)
yield dict(downloaded="[red]FAILED")
if debug_logger:
debug_logger.log(
level="ERROR",
operation="downloader_failed",
message=f"Download failed: {e}",
error=e,
context={
"url_count": len(urls),
"output_dir": str(output_dir),
},
)
raise
yield event_queue.get_nowait()
except Empty:
break
if debug_logger:
debug_logger.log(