loop-uh
c015fbb376
Some checks failed
Windows project source guards / test (push) Failing after 26s
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
247 lines
9.8 KiB
Python
247 lines
9.8 KiB
Python
"""Скачивание файлов и заплаток из хранилища: с докачкой и запасными источниками.
|
||
|
||
Имя файла в хранилище — контрольная сумма содержимого, поэтому у всех
|
||
источников (Forgejo, зеркала) байты одни и те же: оборванный кусок докачивается
|
||
по ``Range`` с любого из них, в том числе после перезапуска программы.
|
||
|
||
Готовый файл распаковывается (заплатка — накладывается на файл с диска),
|
||
проверяется по контрольной сумме из подписанного списка и кладётся в
|
||
``<рабочая папка>/ready/<sha256>``. Файл, уже лежащий там с верной суммой,
|
||
второй раз не качается.
|
||
|
||
Только стандартная библиотека: так одинаково в обеих программах.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import ssl
|
||
import threading
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
from collections.abc import Callable, Sequence
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
|
||
from . import store
|
||
from .manifest import file_sha256, safe_relative_path
|
||
from .plan import UpdatePlan, Want
|
||
|
||
CHUNK = 256 * 1024
|
||
ROUNDS = 3
|
||
ROUND_PAUSES = (2.0, 5.0)
|
||
|
||
|
||
class FetchCancelled(Exception):
|
||
"""Скачивание остановлено по просьбе программы."""
|
||
|
||
|
||
class FetchError(Exception):
|
||
"""Файл не скачался ни с одного источника. ``code`` — причина для сервера."""
|
||
|
||
def __init__(self, message: str, code: str = "net") -> None:
|
||
super().__init__(message)
|
||
self.code = code
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Source:
|
||
"""Один путь к хранилищу: как назвать ссылку на файл по его имени.
|
||
|
||
``proxy`` — адрес прокси (``http://127.0.0.1:порт``) или пусто: напрямую.
|
||
Системный прокси Windows не берётся никогда — путь задаёт программа.
|
||
"""
|
||
|
||
label: str
|
||
url_for: Callable[[str], str]
|
||
verify_ssl: bool = True
|
||
proxy: str = ""
|
||
|
||
|
||
class _Progress:
|
||
def __init__(self, total: int, callback: Callable[[int, int], None]) -> None:
|
||
self.total = max(int(total), 0)
|
||
self.done = 0
|
||
self._callback = callback
|
||
self._lock = threading.Lock()
|
||
|
||
def add(self, amount: int) -> None:
|
||
with self._lock:
|
||
self.done += amount
|
||
done = self.done
|
||
try:
|
||
self._callback(min(done, self.total), self.total)
|
||
except Exception: # noqa: BLE001 — показ хода не должен рвать скачивание
|
||
pass
|
||
|
||
|
||
def _opener(source: Source) -> urllib.request.OpenerDirector:
|
||
context = ssl.create_default_context()
|
||
if not source.verify_ssl:
|
||
context.check_hostname = False
|
||
context.verify_mode = ssl.CERT_NONE
|
||
proxies = {"http": source.proxy, "https": source.proxy} if source.proxy else {}
|
||
return urllib.request.build_opener(urllib.request.ProxyHandler(proxies), urllib.request.HTTPSHandler(context=context))
|
||
|
||
|
||
def _download_part(
|
||
source: Source,
|
||
name: str,
|
||
part: Path,
|
||
size: int,
|
||
*,
|
||
timeout: float,
|
||
cancelled: Callable[[], bool],
|
||
progress: _Progress,
|
||
) -> None:
|
||
"""Докачивает ``part`` до ``size`` байт с одного источника."""
|
||
have = part.stat().st_size if part.exists() else 0
|
||
if have > size:
|
||
part.unlink()
|
||
have = 0
|
||
if have == size:
|
||
return
|
||
headers = {"User-Agent": "ZapretUpdate/1"}
|
||
if have:
|
||
headers["Range"] = f"bytes={have}-"
|
||
request = urllib.request.Request(source.url_for(name), headers=headers)
|
||
with _opener(source).open(request, timeout=timeout) as response:
|
||
status = getattr(response, "status", 200)
|
||
mode = "ab"
|
||
if have and status != 206:
|
||
# Источник не умеет докачку: начинаем этот файл сначала.
|
||
progress.add(-have)
|
||
have, mode = 0, "wb"
|
||
with open(part, mode) as handle:
|
||
while True:
|
||
if cancelled():
|
||
raise FetchCancelled("скачивание остановлено")
|
||
chunk = response.read(CHUNK)
|
||
if not chunk:
|
||
break
|
||
if have + len(chunk) > size:
|
||
raise FetchError(f"{source.label} отдал больше байт, чем в списке выпуска", "size")
|
||
handle.write(chunk)
|
||
have += len(chunk)
|
||
progress.add(len(chunk))
|
||
if have != size:
|
||
raise ConnectionError(f"{source.label}: соединение оборвалось на {have} из {size} байт")
|
||
|
||
|
||
def _finish(want: Want, part: Path, ready: Path, root: Path) -> None:
|
||
data = part.read_bytes()
|
||
entry = want.entry
|
||
if want.patch is None:
|
||
content = store.unpack_blob(data, sha256=entry.sha256, size=entry.size)
|
||
else:
|
||
base = Path(root).joinpath(*safe_relative_path(entry.path).split("/"))
|
||
content = store.apply_patch(
|
||
base.read_bytes(), data, sha256=entry.sha256, size=entry.size, kind=want.patch.kind
|
||
)
|
||
temporary = ready.with_name(ready.name + ".tmp")
|
||
temporary.write_bytes(content)
|
||
os.replace(temporary, ready)
|
||
part.unlink(missing_ok=True)
|
||
|
||
|
||
def _ready_ok(ready: Path, sha256: str, size: int) -> bool:
|
||
try:
|
||
return ready.stat().st_size == size and file_sha256(ready) == sha256
|
||
except OSError:
|
||
return False
|
||
|
||
|
||
def fetch_plan(
|
||
plan: UpdatePlan,
|
||
*,
|
||
root: Path,
|
||
work_dir: Path,
|
||
sources: Sequence[Source],
|
||
threads: int = 4,
|
||
timeout: float = 30.0,
|
||
cancelled: Callable[[], bool] = lambda: False,
|
||
progress: Callable[[int, int], None] = lambda _done, _total: None,
|
||
sleep: Callable[[float], None] = time.sleep,
|
||
) -> dict[str, Path]:
|
||
"""Скачивает всё, что нужно плану; ``{sha256: готовый файл}``."""
|
||
if not sources:
|
||
raise FetchError("нет источников хранилища")
|
||
work = Path(work_dir)
|
||
parts_dir, ready_dir = work / "parts", work / "ready"
|
||
parts_dir.mkdir(parents=True, exist_ok=True)
|
||
ready_dir.mkdir(parents=True, exist_ok=True)
|
||
tracker = _Progress(plan.download_size, progress)
|
||
result: dict[str, Path] = {}
|
||
# Окончательная ошибка одного файла останавливает остальные: без него
|
||
# версию всё равно не собрать.
|
||
failed = threading.Event()
|
||
|
||
def stopped() -> bool:
|
||
return failed.is_set() or cancelled()
|
||
|
||
def one(want: Want) -> None:
|
||
ready = ready_dir / want.entry.sha256
|
||
if _ready_ok(ready, want.entry.sha256, want.entry.size):
|
||
tracker.add(want.download_size)
|
||
result[want.entry.sha256] = ready
|
||
return
|
||
part = parts_dir / (want.name + ".part")
|
||
if part.exists():
|
||
tracker.add(min(part.stat().st_size, want.download_size))
|
||
errors: list[str] = []
|
||
broken = False
|
||
for round_index in range(ROUNDS):
|
||
for source in sources:
|
||
if stopped():
|
||
raise FetchCancelled("скачивание остановлено")
|
||
try:
|
||
_download_part(
|
||
source, want.name, part, want.download_size,
|
||
timeout=timeout, cancelled=stopped, progress=tracker,
|
||
)
|
||
_finish(want, part, ready, root)
|
||
result[want.entry.sha256] = ready
|
||
return
|
||
except FetchCancelled:
|
||
raise
|
||
except store.StoreError as exc:
|
||
# Битый файл: этот кусок не годится ни с какого источника.
|
||
broken = True
|
||
if part.exists():
|
||
tracker.add(-part.stat().st_size)
|
||
part.unlink()
|
||
errors.append(f"{source.label}: {exc}")
|
||
except FetchError as exc:
|
||
if part.exists():
|
||
tracker.add(-part.stat().st_size)
|
||
part.unlink()
|
||
errors.append(f"{source.label}: {exc}")
|
||
except (OSError, urllib.error.URLError, ValueError) as exc:
|
||
errors.append(f"{source.label}: {exc}")
|
||
if round_index < ROUNDS - 1:
|
||
sleep(ROUND_PAUSES[min(round_index, len(ROUND_PAUSES) - 1)])
|
||
raise FetchError(f"{want.entry.path} не скачался: " + "; ".join(errors[-3:]), "hash" if broken else "net")
|
||
|
||
def guarded(want: Want) -> None:
|
||
try:
|
||
one(want)
|
||
except BaseException:
|
||
failed.set()
|
||
raise
|
||
|
||
errors: list[BaseException] = []
|
||
with ThreadPoolExecutor(max_workers=max(1, threads)) as pool:
|
||
for future in [pool.submit(guarded, want) for want in plan.want]:
|
||
try:
|
||
future.result()
|
||
except BaseException as exc: # noqa: BLE001 — собрать и отдать главную
|
||
errors.append(exc)
|
||
if errors:
|
||
# Настоящая причина важнее «остановлено», которое она вызвала у соседей.
|
||
raise next((exc for exc in errors if not isinstance(exc, FetchCancelled)), errors[0])
|
||
return result
|
||
|
||
|
||
__all__ = ["FetchCancelled", "FetchError", "Source", "fetch_plan"]
|