Update hunt

.gitignore

@@ -6,4 +6,5 @@ merge_dev.bat
 __pycache__
 *.png
 *.txt
-db.sqlite
+db.sqlite
+.?

Client/Scripts/hunt.py

@@ -107,17 +107,16 @@ def upsert_urls(conn: sqlite3.Connection, rows: Sequence[Tuple[str, str]]) -> No
         rows,
     )
 
-def fetch_api_hunt_results(
+def iter_api_hunt_results(
     post_ids: Sequence[str],
     phpsessid: str,
     desc: str,
     stop_event: threading.Event = stop_event,
-) -> Dict[str, str]:
+) -> Iterator[Tuple[str, str]]:
     post_ids_list = list(post_ids)
     if not post_ids_list:
-        return {}
+        return
     payload = {"post_ids": post_ids_list, "phpsessid": phpsessid}
-    results: Dict[str, str] = {}
     try:
         with requests.post(
             f"{ENDPOINT}/pixif_stream",
@@ -156,10 +155,10 @@ def fetch_api_hunt_results(
                                 url = ""
                             if DRY_RUN and url:
                                 tqdm.write(f"{post_id} -> {decode_if_binary(url)}")
-                            results[post_id] = url
                             pbar.update(1)
                             last_progress = now
                             consecutive_timeouts = 0
+                            yield post_id, url
                         if idle_break:
                             break
                         break
@@ -171,7 +170,6 @@ def fetch_api_hunt_results(
     except KeyboardInterrupt:
         stop_event.set()
         raise
-    return results
 
 conn = open_db(DB_PATH)
 
@@ -240,44 +238,47 @@ def download_one_image(
                 pass
         return post_id, "error", str(exc)
 
-def download_images(
-    to_download: Dict[str, str],
-    dest_dir: str,
-    phpsessid: str,
-    max_workers: int = MAX_WORKERS,
+def link_group_image(post_id: str, group_dir: str, post_indexes: Dict[str, int]) -> None:
+    index = post_indexes.get(post_id)
+    if index is None:
+        return
+    stash_path = os.path.join(STASH_DIR, f"{post_id}.png")
+    dest_path = os.path.join(group_dir, f"{index}_{post_id}.png")
+    if os.path.exists(stash_path) and not os.path.exists(dest_path):
+        os.link(stash_path, dest_path)
+
+def handle_download_result(future, group_dir: str, post_indexes: Dict[str, int], pbar: tqdm) -> None:
+    post_id, status, detail = future.result()
+    if status == "error":
+        tqdm.write(f"Failed {post_id}: {detail}")
+    elif status in {"ok", "exists"}:
+        images_cache.add(f"{post_id}.png")
+        link_group_image(post_id, group_dir, post_indexes)
+    pbar.update(1)
+
+def drain_downloads(
+    futures,
+    group_dir: str,
+    post_indexes: Dict[str, int],
+    pbar: tqdm,
+) -> None:
+    for future in list(futures):
+        if future.done():
+            futures.remove(future)
+            handle_download_result(future, group_dir, post_indexes, pbar)
+
+def finish_downloads(
+    futures,
+    group_dir: str,
+    post_indexes: Dict[str, int],
+    pbar: tqdm,
     stop_event: threading.Event = stop_event,
 ) -> None:
-    if not to_download:
-        return
-
-    os.makedirs(dest_dir, exist_ok=True)
-    futures = []
-    executor = ThreadPoolExecutor(max_workers=max_workers)
-    interrupted = False
-    try:
-        for post_id, url in to_download.items():
-            if stop_event.is_set():
-                break
-            futures.append(executor.submit(download_one_image, post_id, url, dest_dir, phpsessid, stop_event))
-
-        with tqdm(total=len(futures), unit="image", desc="Downloading images") as pbar:
-            for future in as_completed(futures):
-                if stop_event.is_set():
-                    break
-                post_id, status, detail = future.result()
-                if status == "error":
-                    tqdm.write(f"Failed {post_id}: {detail}")
-                pbar.update(1)
-    except KeyboardInterrupt:
-        interrupted = True
-        stop_event.set()
-        for future in futures:
-            future.cancel()
-        executor.shutdown(wait=False, cancel_futures=True)
-        raise
-    finally:
-        if not interrupted:
-            executor.shutdown(wait=True)
+    for future in as_completed(list(futures)):
+        if stop_event.is_set():
+            break
+        futures.remove(future)
+        handle_download_result(future, group_dir, post_indexes, pbar)
 
 def decode_if_binary(val: Union[str, bytes]) -> str:
     if type(val) is bytes:
@@ -289,6 +290,12 @@ def scan_with_retries(
     phpsessid: str,
     conn: sqlite3.Connection,
     post_ids_dict: Dict[str, Optional[str]],
+    executor: ThreadPoolExecutor,
+    futures,
+    queued_downloads,
+    group_dir: str,
+    post_indexes: Dict[str, int],
+    download_pbar: tqdm,
     desc: str,
     stop_event: threading.Event = stop_event,
 ) -> None:
@@ -297,18 +304,26 @@ def scan_with_retries(
     remaining = list(post_ids)
     attempts = 0
     while remaining and attempts < STREAM_MAX_RETRIES and not stop_event.is_set():
-        results = fetch_api_hunt_results(remaining, phpsessid, desc, stop_event)
-        post_ids_dict.update(results)
-        remaining = [post_id for post_id in remaining if post_id not in results]
+        received = set()
+        for post_id, url in iter_api_hunt_results(remaining, phpsessid, desc, stop_event):
+            received.add(post_id)
+            post_ids_dict[post_id] = url
+            if not DRY_RUN:
+                with conn:
+                    upsert_urls(conn, [(post_id, url)])
+            if url and f"{post_id}.png" not in images_cache and post_id not in queued_downloads:
+                queued_downloads.add(post_id)
+                future = executor.submit(download_one_image, post_id, url, STASH_DIR, phpsessid, stop_event)
+                futures.append(future)
+                download_pbar.total += 1
+                download_pbar.refresh()
+            drain_downloads(futures, group_dir, post_indexes, download_pbar)
+        remaining = [post_id for post_id in remaining if post_id not in received]
         if not remaining:
             break
         attempts += 1
         if STREAM_RETRY_DELAY_SECONDS:
             time.sleep(STREAM_RETRY_DELAY_SECONDS)
-    if not DRY_RUN:
-        rows = [(post_id, url) for post_id, url in post_ids_dict.items() if url is not None]
-        with conn:
-            upsert_urls(conn, rows)
 
 try:
     for index in indexs:
@@ -318,6 +333,7 @@ try:
         with open(os.path.join(ROOT_DIR, valid[index]), "r") as f:
             post_ids = f.read().split()
 
+        post_indexes = {post_id: i for i, post_id in enumerate(post_ids)}
         post_ids_dict = fetch_cached_state(conn, post_ids)
         if DRY_RUN:
             filtered = list(post_ids)
@@ -329,28 +345,60 @@ try:
             ]
         print(f"Group: {group_name}\nFiltered: {len(filtered)}/{len(post_ids)}")
 
-        if filtered:
-            if DRY_RUN:
-                print("Dry run outputs (post_id -> page):")
-            scan_with_retries(filtered, phpsessid, conn, post_ids_dict, "API hunt", stop_event)
         if DRY_RUN:
+            if filtered:
+                print("Dry run outputs (post_id -> page):")
+                for _ in iter_api_hunt_results(filtered, phpsessid, "API hunt", stop_event):
+                    pass
             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}
 
-        if to_download:
-            print(f"Total images to download: {len(to_download)}")
-            if not DRY_RUN:
-                download_images(to_download, STASH_DIR, phpsessid, stop_event=stop_event)
-
-        images_cache.update(os.listdir(STASH_DIR))
+        cached_downloads = {
+            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
+        }
+        futures = []
+        queued_downloads = set()
+        executor = ThreadPoolExecutor(max_workers=MAX_WORKERS)
+        interrupted = False
+        try:
+            with tqdm(total=0, unit="image", desc="Downloading images") as download_pbar:
+                for post_id, url in cached_downloads.items():
+                    queued_downloads.add(post_id)
+                    futures.append(executor.submit(download_one_image, post_id, url, STASH_DIR, phpsessid, stop_event))
+                    download_pbar.total += 1
+                    download_pbar.refresh()
+                if filtered:
+                    scan_with_retries(
+                        filtered,
+                        phpsessid,
+                        conn,
+                        post_ids_dict,
+                        executor,
+                        futures,
+                        queued_downloads,
+                        group_dir,
+                        post_indexes,
+                        download_pbar,
+                        "API hunt",
+                        stop_event,
+                    )
+                finish_downloads(futures, group_dir, post_indexes, download_pbar, stop_event)
+        except KeyboardInterrupt:
+            interrupted = True
+            stop_event.set()
+            for future in futures:
+                future.cancel()
+            executor.shutdown(wait=False, cancel_futures=True)
+            raise
+        finally:
+            if not interrupted:
+                executor.shutdown(wait=True)
 
         print("Linking images to the group directory...")
-        for i, post_id in enumerate(post_ids):
-            stash_path = os.path.join(STASH_DIR, f"{post_id}.png")
-            dest_path = os.path.join(group_dir, f"{i}_{post_id}.png")
-
-            if os.path.exists(stash_path) and not os.path.exists(dest_path):
-                os.link(stash_path, dest_path)
+        images_cache.update(os.listdir(STASH_DIR))
+        for post_id in post_ids:
+            link_group_image(post_id, group_dir, post_indexes)
 
         if len(os.listdir(group_dir)) == 0:
             shutil.rmtree(group_dir)