Retry rate-limited HLS segments

provider_downloader.py

@@ -9,6 +9,7 @@ import re
 import shutil
 import subprocess
 import sys
+import time
 from urllib.parse import urljoin, urlparse
 import urllib.request
 from pathlib import Path
@@ -17,6 +18,7 @@ from pathlib import Path
 PROJECT_ROOT = Path(__file__).resolve().parent
 BRIDGE = PROJECT_ROOT / "providers" / "bridge.js"
 PROVIDERS = ("anikoto", "anineko", "pahe")
+RETRYABLE_HTTP_STATUS = {408, 425, 429, 500, 502, 503, 504}
 
 
 def clean_component(value, default="Anime"):
@@ -218,6 +220,61 @@ def fetch_bytes(url, headers=None, timeout=30):
         return response.read()
 
 
+def int_env(name, default, minimum=0):
+    try:
+        value = int(str(os.environ.get(name, default)).strip())
+    except (TypeError, ValueError):
+        return default
+    return max(minimum, value)
+
+
+def float_env(name, default, minimum=0.0):
+    try:
+        value = float(str(os.environ.get(name, default)).strip())
+    except (TypeError, ValueError):
+        return default
+    return max(minimum, value)
+
+
+def retry_after_seconds(exc):
+    try:
+        value = exc.headers.get("Retry-After")
+    except Exception:
+        return None
+    if not value:
+        return None
+    text = str(value).strip()
+    try:
+        return max(0.0, float(text))
+    except ValueError:
+        return None
+
+
+def fetch_bytes_with_retries(url, headers=None, timeout=30, retries=None, retry_label="request"):
+    attempts = int_env("KAIZOKU_SEGMENT_RETRIES", 8, minimum=0) if retries is None else max(0, int(retries))
+    base_delay = float_env("KAIZOKU_SEGMENT_RETRY_DELAY", 1.25, minimum=0.0)
+    max_delay = float_env("KAIZOKU_SEGMENT_RETRY_MAX_DELAY", 15.0, minimum=0.0)
+    for attempt in range(attempts + 1):
+        try:
+            return fetch_bytes(url, headers=headers, timeout=timeout)
+        except urllib.error.HTTPError as exc:
+            retryable = exc.code in RETRYABLE_HTTP_STATUS
+            if not retryable or attempt >= attempts:
+                raise
+            header_delay = retry_after_seconds(exc)
+            delay = header_delay if header_delay is not None else base_delay * (2 ** attempt)
+            if max_delay:
+                delay = min(delay, max_delay)
+            print(
+                f"{retry_label} got HTTP {exc.code}; retrying in {delay:.1f}s "
+                f"({attempt + 1}/{attempts}).",
+                file=sys.stderr,
+                flush=True,
+            )
+            if delay:
+                time.sleep(delay)
+
+
 def fetch_text(url, headers=None, timeout=30):
     return fetch_bytes(url, headers=headers, timeout=timeout).decode("utf-8", errors="replace")
 
@@ -393,12 +450,22 @@ def download_hls_segments(stream, input_url, target, partial, episode_number=Non
         last_percent = -1
         with ts_file.open("wb") as joined:
             for index, segment in enumerate(segments, start=1):
-                data = fetch_bytes(segment["url"], headers=stream.get("headers") or {}, timeout=60)
+                data = fetch_bytes_with_retries(
+                    segment["url"],
+                    headers=stream.get("headers") or {},
+                    timeout=60,
+                    retry_label=f"Episode {episode_number} segment {index}/{len(segments)}",
+                )
                 data = strip_png_header(data)
                 key_url = segment.get("key_url")
                 if key_url:
                     if key_url not in key_cache:
-                        key_cache[key_url] = fetch_bytes(key_url, headers=stream.get("headers") or {}, timeout=30)
+                        key_cache[key_url] = fetch_bytes_with_retries(
+                            key_url,
+                            headers=stream.get("headers") or {},
+                            timeout=30,
+                            retry_label=f"Episode {episode_number} segment key",
+                        )
                     data = decrypt_aes128_segment(data, key_cache[key_url], segment.get("iv"))
                 if not data:
                     raise RuntimeError(f"Segment {index} was empty.")