Optimize scanning and image proxy paths

backend/app.py

@@ -77,6 +77,7 @@ THUMB_MAX_AGE = 1800
 THUMB_CLEAN_INTERVAL = 300
 THUMB_DIR = Path(tempfile.gettempdir()) / "pixif2-thumbs"
 PAGE_URL_CACHE_MAX_AGE = 1800
+PAGE_URL_CACHE_LIMIT = 5000
 WEBP_SCALE = 0.4
 WEBP_QUALITY = 82
 STEALTH_ALPHA_SIGS = (b"stealth_pnginfo", b"stealth_pngcomp")
@@ -86,7 +87,9 @@ app = FastAPI()
 ACTIVE_TASKS = {}
 TASK_EVENT_QUEUES = set()
 PAGE_URL_CACHE = {}
+PAGE_URL_INFLIGHT = {}
 TURSO_CLIENT = None
+PIXIV_CLIENT = None
 LAST_THUMB_CLEANUP = 0
 
 FRONTEND_DIR = os.path.join(os.getcwd(), "frontend")
@@ -132,6 +135,15 @@ async def turso_client():
     return TURSO_CLIENT
 
 
+async def pixiv_client():
+    global PIXIV_CLIENT
+    if PIXIV_CLIENT is None or PIXIV_CLIENT.closed:
+        PIXIV_CLIENT = aiohttp.ClientSession(
+            cookies={"PHPSESSID": PHPSESSID}, headers=PIXIV_HEADERS
+        )
+    return PIXIV_CLIENT
+
+
 async def turso_execute(stmts):
     url = turso_url() + "/v2/pipeline"
     headers = {
@@ -685,20 +697,40 @@ async def get_post_pages(post_id, session):
     now = time.time()
     if cached and now - cached["time"] < PAGE_URL_CACHE_MAX_AGE:
         return cached["pages"]
-    data = await fetch_page(session, f"https://www.pixiv.net/ajax/illust/{post_id}/pages")
-    pages = []
-    for page in data.get("body") or []:
-        urls = page.get("urls") or {}
-        pages.append(
-            {
-                "original": urls.get("original") or "",
-                "regular": urls.get("regular") or "",
-                "small": urls.get("small") or "",
-                "thumb_mini": urls.get("thumb_mini") or "",
-            }
+    inflight = PAGE_URL_INFLIGHT.get(key)
+    if inflight:
+        return await inflight
+
+    async def load():
+        data = await fetch_page(
+            session, f"https://www.pixiv.net/ajax/illust/{post_id}/pages"
         )
-    PAGE_URL_CACHE[key] = {"time": now, "pages": pages}
-    return pages
+        pages = []
+        for page in data.get("body") or []:
+            urls = page.get("urls") or {}
+            pages.append(
+                {
+                    "original": urls.get("original") or "",
+                    "regular": urls.get("regular") or "",
+                    "small": urls.get("small") or "",
+                    "thumb_mini": urls.get("thumb_mini") or "",
+                }
+            )
+        PAGE_URL_CACHE[key] = {"time": time.time(), "pages": pages}
+        if len(PAGE_URL_CACHE) > PAGE_URL_CACHE_LIMIT:
+            cutoff = time.time() - PAGE_URL_CACHE_MAX_AGE
+            for cid, row in list(PAGE_URL_CACHE.items()):
+                if row["time"] < cutoff:
+                    PAGE_URL_CACHE.pop(cid, None)
+        return pages
+
+    task = asyncio.create_task(load())
+    PAGE_URL_INFLIGHT[key] = task
+    try:
+        return await task
+    finally:
+        if PAGE_URL_INFLIGHT.get(key) is task:
+            PAGE_URL_INFLIGHT.pop(key, None)
 
 
 def decode_meta_text(data):
@@ -829,15 +861,22 @@ def parse_xmp_metadata(data):
 
 def parse_image_metadata(data):
     fields = {}
-    for parser in (
-        parse_png_metadata,
-        parse_pil_metadata,
-        parse_jpeg_comments,
-        parse_xmp_metadata,
-    ):
+    if data.startswith(b"\x89PNG\r\n\x1a\n"):
+        parsers = (parse_png_metadata, parse_xmp_metadata)
+    elif data.startswith(b"\xff\xd8"):
+        parsers = (parse_pil_metadata, parse_jpeg_comments, parse_xmp_metadata)
+    elif data.startswith(b"RIFF") and data[8:12] == b"WEBP":
+        parsers = (parse_pil_metadata, parse_xmp_metadata)
+    else:
+        parsers = (parse_pil_metadata, parse_xmp_metadata)
+    for parser in parsers:
         parsed = parser(data)
         if parsed:
             fields.update(parsed)
+    if not fields and data.startswith(b"\x89PNG\r\n\x1a\n"):
+        parsed = parse_pil_metadata(data)
+        if parsed:
+            fields.update(parsed)
     return fields or None
 
 
@@ -994,6 +1033,10 @@ def determine_exif_type(fields):
     return None
 
 
+def image_ext(url):
+    return urlsplit(url).path.rsplit(".", 1)[-1].casefold()
+
+
 def byteize(bits):
     bits = bits[: (bits.shape[0] // 8) * 8]
     bits = np.bitwise_and(bits, 1)
@@ -1026,8 +1069,7 @@ async def scan_post(post_id, session, post_sem, exif_sem, img_sem):
             image_urls = [
                 p["original"]
                 for p in pages
-                if urlsplit(p["original"]).path.rsplit(".", 1)[-1].casefold()
-                in ("png", "jpg", "jpeg", "webp")
+                if image_ext(p["original"]) in ("png", "jpg", "jpeg", "webp")
             ]
             for url in image_urls:
                 metadata = await get_exif_range(url, session, exif_sem)
@@ -1035,7 +1077,7 @@ async def scan_post(post_id, session, post_sem, exif_sem, img_sem):
                 if exif_type:
                     code = EXIF_TYPE_TO_CODE.get(exif_type)
                     return post_id, url, code
-            for url in image_urls:
+            for url in (u for u in image_urls if image_ext(u) == "png"):
                 img_data = await fetch_image(session, url, img_sem)
                 if img_data and has_stealth_png_bytes(img_data):
                     return post_id, url, EXIF_TYPE_TO_CODE.get("stealth")
@@ -1075,25 +1117,42 @@ async def run_scan(post_ids, phpsessid, task_id=None, save_live=False, done_offs
     cookies = {"PHPSESSID": phpsessid}
     results = []
     async with aiohttp.ClientSession(cookies=cookies, headers=PIXIV_HEADERS) as session:
-        tasks = [
-            asyncio.create_task(scan_post(pid, session, post_sem, exif_sem, img_sem))
-            for pid in post_ids
-        ]
+        post_iter = iter(post_ids)
+        tasks = set()
+
+        def fill_tasks():
+            while len(tasks) < POST_SCAN_LIMIT:
+                try:
+                    pid = next(post_iter)
+                except StopIteration:
+                    break
+                tasks.add(
+                    asyncio.create_task(
+                        scan_post(pid, session, post_sem, exif_sem, img_sem)
+                    )
+                )
+
+        fill_tasks()
         pending = []
         try:
-            for coro in asyncio.as_completed(tasks):
+            while tasks:
                 if task_canceled(task_id):
                     raise TaskCanceled()
-                result = await coro
-                results.append(result)
-                pending.append(result)
-                if task_id and task_id in ACTIVE_TASKS:
-                    ACTIVE_TASKS[task_id]["done"] = done_offset + len(results)
-                if save_live and len(pending) >= 20:
-                    if task_canceled(task_id):
-                        raise TaskCanceled()
-                    await save_scan_results(pending)
-                    pending = []
+                done, tasks = await asyncio.wait(
+                    tasks, return_when=asyncio.FIRST_COMPLETED
+                )
+                for task in done:
+                    result = task.result()
+                    results.append(result)
+                    pending.append(result)
+                    if task_id and task_id in ACTIVE_TASKS:
+                        ACTIVE_TASKS[task_id]["done"] = done_offset + len(results)
+                    if save_live and len(pending) >= 20:
+                        if task_canceled(task_id):
+                            raise TaskCanceled()
+                        await save_scan_results(pending)
+                        pending = []
+                fill_tasks()
             if save_live and pending:
                 if task_canceled(task_id):
                     raise TaskCanceled()
@@ -1289,9 +1348,8 @@ def media_type_from_url(url):
 
 
 async def get_pixiv_image_url(post_id, page, size, phpsessid):
-    cookies = {"PHPSESSID": phpsessid}
-    async with aiohttp.ClientSession(cookies=cookies, headers=PIXIV_HEADERS) as session:
-        pages = await get_post_pages(post_id, session)
+    session = await pixiv_client()
+    pages = await get_post_pages(post_id, session)
     if not pages:
         raise HTTPException(status_code=404, detail="image not found")
     page = min(max(page, 0), len(pages) - 1)
@@ -1306,12 +1364,21 @@ async def get_pixiv_image_url(post_id, page, size, phpsessid):
 
 
 async def fetch_pixiv_bytes(url, phpsessid):
-    cookies = {"PHPSESSID": phpsessid}
-    async with aiohttp.ClientSession(cookies=cookies, headers=PIXIV_HEADERS) as session:
-        async with session.get(url) as r:
-            if r.status >= 400:
-                raise HTTPException(status_code=r.status, detail="image fetch failed")
-            return await r.read(), r.headers.get("Content-Type") or media_type_from_url(url)
+    session = await pixiv_client()
+    async with session.get(url) as r:
+        if r.status >= 400:
+            raise HTTPException(status_code=r.status, detail="image fetch failed")
+        return await r.read(), r.headers.get("Content-Type") or media_type_from_url(url)
+
+
+def write_webp(data, out):
+    image = Image.open(io.BytesIO(data))
+    image = image.resize(
+        (max(int(image.width * WEBP_SCALE), 1), max(int(image.height * WEBP_SCALE), 1))
+    )
+    if image.mode not in ("RGB", "RGBA"):
+        image = image.convert("RGB")
+    image.save(out, "WEBP", quality=WEBP_QUALITY)
 
 
 async def create_webp(post_id, image_url, phpsessid, page=0, kind="t"):
@@ -1324,13 +1391,7 @@ async def create_webp(post_id, image_url, phpsessid, page=0, kind="t"):
     data, _ = await fetch_pixiv_bytes(image_url, phpsessid)
     if not data:
         raise HTTPException(status_code=404, detail="image not found")
-    image = Image.open(io.BytesIO(data))
-    image = image.resize(
-        (max(int(image.width * WEBP_SCALE), 1), max(int(image.height * WEBP_SCALE), 1))
-    )
-    if image.mode not in ("RGB", "RGBA"):
-        image = image.convert("RGB")
-    image.save(out, "WEBP", quality=WEBP_QUALITY)
+    await asyncio.to_thread(write_webp, data, out)
     return out
 
 
@@ -1577,6 +1638,8 @@ async def startup():
 async def shutdown():
     if TURSO_CLIENT and not TURSO_CLIENT.is_closed:
         await TURSO_CLIENT.aclose()
+    if PIXIV_CLIENT and not PIXIV_CLIENT.closed:
+        await PIXIV_CLIENT.close()
 
 
 async def backfill_search_stats():

frontend/app.js

@@ -434,6 +434,7 @@ async function openSearch(id, page = 1, show = false, view = explorerDetailView(
 }
 
 async function loadReclass(id) {
+  if (!id.endsWith("_reclass")) return null
   try {
     const resp = await fetch(`/api/reclass/${encodeURIComponent(id)}`)
     const data = await resp.json()