Supposedly faster

Client/hunt.py

@@ -2,8 +2,10 @@ import json
 import os
 import sqlite3
 import threading
+import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import shutil
+from typing import Dict, Iterator, List, Optional, Sequence, Tuple, Union
 
 import requests
 from tqdm import tqdm
@@ -18,9 +20,14 @@ REQUEST_HEADERS = {
 }
 MAX_WORKERS = min(16, os.cpu_count() or 8)
 REQUEST_TIMEOUT = (120, 120)
+STREAM_REQUEST_TIMEOUT = (120, 30)
+STREAM_IDLE_TIMEOUT_SECONDS = 45
+STREAM_MAX_READ_TIMEOUTS = 3
+STREAM_MAX_RETRIES = 3
+STREAM_RETRY_DELAY_SECONDS = 2
 stop_event = threading.Event()
 
-def read_dotenv_value(path, key):
+def read_dotenv_value(path: str, key: str) -> Optional[str]:
     try:
         with open(path, "r") as env_file:
             for line in env_file:
@@ -34,7 +41,7 @@ def read_dotenv_value(path, key):
         return None
     return None
 
-def get_phpsessid():
+def get_phpsessid() -> str:
     phpsessid = os.getenv("PHPSESSID")
     if phpsessid:
         return phpsessid
@@ -53,7 +60,7 @@ images_cache = set(os.listdir("images/Stash"))
 
 DB_PATH = "db.sqlite"
 
-def open_db(path):
+def open_db(path: str) -> sqlite3.Connection:
     conn = sqlite3.connect(path)
     conn.execute(
         """
@@ -66,11 +73,11 @@ def open_db(path):
     conn.commit()
     return conn
 
-def chunked(seq, size):
+def chunked(seq: Sequence[str], size: int) -> Iterator[List[str]]:
     for i in range(0, len(seq), size):
         yield seq[i:i + size]
 
-def fetch_cached_urls(conn, post_ids):
+def fetch_cached_urls(conn: sqlite3.Connection, post_ids: Sequence[str]) -> Dict[str, Optional[str]]:
     post_ids_dict = {post_id: None for post_id in post_ids}
     if not post_ids:
         return post_ids_dict
@@ -83,7 +90,7 @@ def fetch_cached_urls(conn, post_ids):
 
     return post_ids_dict
 
-def upsert_urls(conn, rows):
+def upsert_urls(conn: sqlite3.Connection, rows: Sequence[Tuple[str, str]]) -> None:
     if not rows:
         return
     conn.executemany(
@@ -95,24 +102,42 @@ def upsert_urls(conn, rows):
         rows,
     )
 
-def stream_pixif_updates(post_ids, phpsessid, conn, post_ids_dict, desc, stop_event=stop_event):
-    if not post_ids:
-        return
-    payload = {"post_ids": post_ids, "phpsessid": phpsessid}
+def stream_pixif_updates(
+    post_ids: Sequence[str],
+    phpsessid: str,
+    conn: sqlite3.Connection,
+    post_ids_dict: Dict[str, Optional[str]],
+    desc: str,
+    stop_event: threading.Event = stop_event,
+) -> int:
+    post_ids_list = list(post_ids)
+    if not post_ids_list:
+        return 0
+    payload = {"post_ids": post_ids_list, "phpsessid": phpsessid}
+    processed = 0
     try:
         with requests.post(
             f"{endpoint}/pixif_stream",
             json=payload,
             stream=True,
-            timeout=REQUEST_TIMEOUT,
+            timeout=STREAM_REQUEST_TIMEOUT,
         ) as response:
             response.raise_for_status()
-            with tqdm(total=len(post_ids), unit="post", desc=desc) as pbar:
+            with tqdm(total=len(post_ids_list), unit="post", desc=desc) as pbar:
+                last_progress = time.monotonic()
+                consecutive_timeouts = 0
+                idle_break = False
                 while not stop_event.is_set() and pbar.n < pbar.total:
+                    if time.monotonic() - last_progress > STREAM_IDLE_TIMEOUT_SECONDS:
+                        break
                     try:
                         for line in response.iter_lines(decode_unicode=True):
                             if stop_event.is_set():
                                 break
+                            now = time.monotonic()
+                            if now - last_progress > STREAM_IDLE_TIMEOUT_SECONDS:
+                                idle_break = True
+                                break
                             if not line:
                                 continue
                             try:
@@ -133,12 +158,21 @@ def stream_pixif_updates(post_ids, phpsessid, conn, post_ids_dict, desc, stop_ev
                                     upsert_urls(conn, [(post_id, url)])
                             post_ids_dict[post_id] = url
                             pbar.update(1)
+                            processed += 1
+                            last_progress = now
+                            consecutive_timeouts = 0
+                        if idle_break:
+                            break
                         break
                     except requests.exceptions.ReadTimeout:
+                        consecutive_timeouts += 1
+                        if consecutive_timeouts >= STREAM_MAX_READ_TIMEOUTS:
+                            break
                         continue
     except KeyboardInterrupt:
         stop_event.set()
         raise
+    return processed
 
 conn = open_db(DB_PATH)
 
@@ -155,12 +189,18 @@ for inp in inputs:
     elif inp.isdigit():
         indexs.append(int(inp) - 1)
 
-def build_image_url(url):
+def build_image_url(url: str) -> str:
     if url.startswith("http"):
         return url
     return IMG_BASE + url
 
-def download_one_image(post_id, url, dest_dir, phpsessid, stop_event=stop_event):
+def download_one_image(
+    post_id: str,
+    url: str,
+    dest_dir: str,
+    phpsessid: str,
+    stop_event: threading.Event = stop_event,
+) -> Tuple[str, str, Optional[str]]:
     dest_path = os.path.join(dest_dir, f"{post_id}.png")
     if os.path.exists(dest_path):
         return post_id, "exists", None
@@ -201,7 +241,13 @@ def download_one_image(post_id, url, dest_dir, phpsessid, stop_event=stop_event)
                 pass
         return post_id, "error", str(exc)
 
-def download_images(to_download, dest_dir, phpsessid, max_workers=MAX_WORKERS, stop_event=stop_event):
+def download_images(
+    to_download: Dict[str, str],
+    dest_dir: str,
+    phpsessid: str,
+    max_workers: int = MAX_WORKERS,
+    stop_event: threading.Event = stop_event,
+) -> None:
     if not to_download:
         return
 
@@ -234,11 +280,32 @@ def download_images(to_download, dest_dir, phpsessid, max_workers=MAX_WORKERS, s
         if not interrupted:
             executor.shutdown(wait=True)
 
-def decode_if_binary(val):
+def decode_if_binary(val: Union[str, bytes]) -> str:
     if type(val) is bytes:
         return val.decode()
     return val
 
+def scan_with_retries(
+    post_ids: Sequence[str],
+    phpsessid: str,
+    conn: sqlite3.Connection,
+    post_ids_dict: Dict[str, Optional[str]],
+    desc: str,
+    stop_event: threading.Event = stop_event,
+) -> None:
+    if not post_ids:
+        return
+    remaining = list(post_ids)
+    attempts = 0
+    while remaining and attempts < STREAM_MAX_RETRIES and not stop_event.is_set():
+        stream_pixif_updates(remaining, phpsessid, conn, post_ids_dict, desc, stop_event)
+        remaining = [post_id for post_id in remaining if post_ids_dict.get(post_id) is None]
+        if not remaining:
+            break
+        attempts += 1
+        if STREAM_RETRY_DELAY_SECONDS:
+            time.sleep(STREAM_RETRY_DELAY_SECONDS)
+
 try:
     for index in indexs:
         group_name = valid[index].rsplit(".", 1)[0]
@@ -260,7 +327,7 @@ try:
         if filtered:
             if dry_run:
                 print("Dry run outputs (post_id -> page):")
-            stream_pixif_updates(filtered, phpsessid, conn, post_ids_dict, "Scanning", stop_event)
+            scan_with_retries(filtered, phpsessid, conn, post_ids_dict, "Scanning", stop_event)
         if dry_run:
             continue
         to_download = {post_id: decode_if_binary(url) for post_id, url in post_ids_dict.items() if url and f"{post_id}.png" not in images_cache}