refactor: 重构缓存相关逻辑，以及触发额外搜图条件时先发送之前的搜图结果，提高响应速度

Author: NekoAria

YetAnotherPicSearch/__init__.py

@@ -1,13 +1,20 @@
 import asyncio
 import re
-from asyncio import Lock
 from collections import defaultdict
 from contextlib import suppress
-from typing import Any, DefaultDict, Dict, List, Optional, Tuple, Union
+from typing import (
+    Any,
+    DefaultDict,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
 
 import arrow
 from aiohttp import ClientSession
-from diskcache import Cache
+from cachetools import TTLCache
 from nonebot.adapters.onebot.v11 import (
     ActionFailed,
     Bot,
@@ -22,20 +29,31 @@
 from nonebot.plugin.on import on_message, on_metaevent
 from nonebot.rule import Rule
 from PicImageSearch import Network
+from shelved_cache import PersistentCache
 from tenacity import retry, stop_after_attempt, stop_after_delay
 
 from .ascii2d import ascii2d_search
 from .baidu import baidu_search
-from .cache import exist_in_cache, upsert_cache
 from .config import config
 from .ehentai import ehentai_search
 from .iqdb import iqdb_search
 from .saucenao import saucenao_search
-from .utils import DEFAULT_HEADERS, get_bot_friend_list, handle_img, handle_reply_msg
+from .utils import (
+    DEFAULT_HEADERS,
+    SEARCH_FUNCTION_TYPE,
+    get_bot_friend_list,
+    handle_reply_msg,
+)
 
 sending_lock: DefaultDict[Tuple[Union[int, str], str], asyncio.Lock] = defaultdict(
     asyncio.Lock
 )
+pic_search_cache = PersistentCache(
+    TTLCache,
+    filename="pic_search_cache",
+    maxsize=config.cache_expire * 100,
+    ttl=config.cache_expire * 24 * 60 * 60,
+)
 
 
 def check_first_connect(_: LifecycleMetaEvent) -> bool:
@@ -91,42 +109,58 @@ async def handle_first_receive(event: MessageEvent, matcher: Matcher) -> None:
 
 
 async def image_search(
+    bot: Bot,
+    event: MessageEvent,
     url: str,
     md5: str,
     mode: str,
     purge: bool,
-    _cache: Cache,
     client: ClientSession,
-) -> List[str]:
+    index: Optional[int] = None,
+) -> None:
     url = await get_universal_img_url(url)
-    if not purge and (result := exist_in_cache(_cache, md5, mode)):
-        return [f"[缓存] {i}" for i in result]
+    cache_key = f"{md5}_{mode}"
     try:
-        result = await handle_search_mode(url, md5, mode, _cache, client)
+        if not purge and cache_key in pic_search_cache:
+            result, extra_handle = pic_search_cache[cache_key]
+            await send_result_message(bot, event, [f"[缓存] {i}" for i in result], index)
+            if callable(extra_handle):
+                await send_result_message(
+                    bot, event, await extra_handle(url, client), index
+                )
+            return
+
+        result, extra_handle = await handle_search_mode(url, md5, mode, client)
+        await send_result_message(bot, event, result, index)
+        if callable(extra_handle):
+            await send_result_message(
+                bot, event, await extra_handle(url, client), index
+            )
     except Exception as e:
         logger.exception(f"该图 [{url}] 搜图失败")
-        result = [f"该图搜图失败\nE: {repr(e)}"]
-    return result
+        await send_result_message(bot, event, [f"该图搜图失败\nE: {repr(e)}"], index)
 
 
 @retry(stop=(stop_after_attempt(3) | stop_after_delay(30)), reraise=True)
 async def handle_search_mode(
-    url: str, md5: str, mode: str, _cache: Cache, client: ClientSession
-) -> List[str]:
+    url: str, md5: str, mode: str, client: ClientSession
+) -> Tuple[List[str], Optional[SEARCH_FUNCTION_TYPE]]:
+    extra_handle = None
+
     if mode == "a2d":
         result = await ascii2d_search(url, client)
     elif mode == "ex":
-        result = await ehentai_search(url, client)
+        result, extra_handle = await ehentai_search(url, client)
     elif mode == "iqdb":
-        result = await iqdb_search(url, client)
+        result, extra_handle = await iqdb_search(url, client)
     elif mode == "baidu":
         result = await baidu_search(url, client)
     else:
-        result = await saucenao_search(url, client, mode)
-        # 仅对涉及到 saucenao 的搜图结果做缓存
-        upsert_cache(_cache, md5, mode, result)
+        result, extra_handle = await saucenao_search(url, client, mode)
 
-    return result
+    pic_search_cache[f"{md5}_{mode}"] = result, extra_handle
+
+    return result, extra_handle
 
 
 async def get_universal_img_url(url: str) -> str:
@@ -196,7 +230,7 @@ async def send_message_with_lock(
     bot: Bot,
     event: MessageEvent,
     msg_list: List[str],
-    current_sending_lock: Lock,
+    current_sending_lock: asyncio.Lock,
     index: Optional[int] = None,
 ) -> None:
     start_time = arrow.now()
@@ -269,14 +303,16 @@ async def handle_image_search(bot: Bot, event: MessageEvent, matcher: Matcher) -
         else Network(proxies=config.proxy)
     )
     async with network as client:
-        with Cache("picsearch_cache") as _cache:
-            for index, (url, md5) in enumerate(image_urls_with_md5):
-                await send_result_message(
-                    bot,
-                    event,
-                    await image_search(url, md5, mode, purge, _cache, client),
-                    index if len(image_urls_with_md5) > 1 else None,
-                )
-            _cache.expire()
+        for index, (url, md5) in enumerate(image_urls_with_md5):
+            await image_search(
+                bot,
+                event,
+                url,
+                md5,
+                mode,
+                purge,
+                client,
+                index if len(image_urls_with_md5) > 1 else None,
+            )
 
     await bot.delete_msg(message_id=searching_tips["message_id"])

YetAnotherPicSearch/cache.py

@@ -3,7 +3,7 @@
 from asyncio import sleep
 from collections import defaultdict
 from difflib import SequenceMatcher
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional, Tuple
 
 import arrow
 from aiohttp import ClientSession
@@ -13,7 +13,7 @@
 
 from .ascii2d import ascii2d_search
 from .config import config
-from .utils import DEFAULT_HEADERS, handle_img, shorten_url
+from .utils import DEFAULT_HEADERS, SEARCH_FUNCTION_TYPE, handle_img, shorten_url
 
 EHENTAI_HEADERS = (
     {"Cookie": config.exhentai_cookies, **DEFAULT_HEADERS}
@@ -22,7 +22,9 @@
 )
 
 
-async def ehentai_search(url: str, client: ClientSession) -> List[str]:
+async def ehentai_search(
+    url: str, client: ClientSession
+) -> Tuple[List[str], Optional[SEARCH_FUNCTION_TYPE]]:
     ex = bool(config.exhentai_cookies)
     ehentai = EHentai(client=client)
 
@@ -40,10 +42,11 @@ async def ehentai_search(url: str, client: ClientSession) -> List[str]:
 
         if not res.raw and config.auto_use_ascii2d:
             final_res.append("自动使用 Ascii2D 进行搜索")
-            final_res.extend(await ascii2d_search(url, client))
-        return final_res
+            return final_res, ascii2d_search
 
-    return ["EHentai 暂时无法使用"]
+        return final_res, None
+
+    return ["EHentai 暂时无法使用"], None
 
 
 async def ehentai_title_search(title: str) -> List[str]:

YetAnotherPicSearch/ehentai.py

@@ -1,19 +1,21 @@
-from typing import List
+from typing import List, Optional, Tuple
 
 from aiohttp import ClientSession
 from PicImageSearch import Iqdb
 from yarl import URL
 
 from .ascii2d import ascii2d_search
 from .config import config
-from .utils import get_source, handle_img, shorten_url
+from .utils import SEARCH_FUNCTION_TYPE, get_source, handle_img, shorten_url
 
 
-async def iqdb_search(url: str, client: ClientSession) -> List[str]:
+async def iqdb_search(
+    url: str, client: ClientSession
+) -> Tuple[List[str], Optional[SEARCH_FUNCTION_TYPE]]:
     iqdb = Iqdb(client=client)
     res = await iqdb.search(url)
     if not res.raw:
-        return ["Iqdb 暂时无法使用"]
+        return ["Iqdb 暂时无法使用，自动使用 Ascii2D 进行搜索"], ascii2d_search
 
     final_res: List[str] = []
     # 如果遇到搜索结果相似度低的情况，去除第一个只有提示信息的空结果
@@ -49,6 +51,6 @@ async def iqdb_search(url: str, client: ClientSession) -> List[str]:
 
     if low_acc and config.auto_use_ascii2d:
         final_res.append(f"相似度 {selected_res.similarity}% 过低，自动使用 Ascii2D 进行搜索")
-        final_res.extend(await ascii2d_search(url, client))
+        return final_res, ascii2d_search
 
-    return final_res
+    return final_res, None

YetAnotherPicSearch/iqdb.py

@@ -1,6 +1,6 @@
 import re
 from asyncio import sleep
-from typing import List
+from typing import List, Optional, Tuple
 
 from aiohttp import ClientSession
 from PicImageSearch import SauceNAO
@@ -10,7 +10,7 @@
 from .ascii2d import ascii2d_search
 from .config import config
 from .ehentai import ehentai_title_search
-from .utils import get_source, handle_img, shorten_url
+from .utils import SEARCH_FUNCTION_TYPE, get_source, handle_img, shorten_url
 from .whatanime import whatanime_search
 
 SAUCENAO_DB = {
@@ -23,7 +23,9 @@
 }
 
 
-async def saucenao_search(url: str, client: ClientSession, mode: str) -> List[str]:
+async def saucenao_search(
+    url: str, client: ClientSession, mode: str
+) -> Tuple[List[str], Optional[SEARCH_FUNCTION_TYPE]]:
     db = SAUCENAO_DB[mode]
     if isinstance(db, list):
         saucenao = SauceNAO(
@@ -51,11 +53,10 @@ async def saucenao_search(url: str, client: ClientSession, mode: str) -> List[st
 
     if not res or not res.raw:
         final_res = ["SauceNAO 暂时无法使用，自动使用 Ascii2D 进行搜索"]
-        final_res.extend(await ascii2d_search(url, client))
-        return final_res
+        return final_res, ascii2d_search
 
     selected_res = get_best_result(res, res.raw[0])
-    return await get_final_res(url, client, mode, res, selected_res)
+    return await get_final_res(mode, res, selected_res)
 
 
 def get_best_pixiv_result(
@@ -90,12 +91,8 @@ def get_best_result(res: SauceNAOResponse, selected_res: SauceNAOItem) -> SauceN
 
 
 async def get_final_res(
-    url: str,
-    client: ClientSession,
-    mode: str,
-    res: SauceNAOResponse,
-    selected_res: SauceNAOItem,
-) -> List[str]:
+    mode: str, res: SauceNAOResponse, selected_res: SauceNAOItem
+) -> Tuple[List[str], Optional[SEARCH_FUNCTION_TYPE]]:
     low_acc = selected_res.similarity < config.saucenao_low_acc
     hide_img = config.hide_img or (
         selected_res.hidden or low_acc and config.hide_img_when_low_acc
@@ -133,7 +130,9 @@ async def get_final_res(
     final_res.append("\n".join([i for i in res_list if i]))
 
     if low_acc:
-        final_res.extend(await handle_saucenao_low_acc(url, client, mode, selected_res))
+        extra_res, extra_handle = await handle_saucenao_low_acc(mode, selected_res)
+        final_res.extend(extra_res)
+        return final_res, extra_handle
     elif selected_res.index_id in SAUCENAO_DB["doujin"]:  # type: ignore
         title = selected_res.title.replace("-", "")
         final_res.extend(await ehentai_title_search(title))
@@ -143,23 +142,20 @@ async def get_final_res(
             await ehentai_title_search(f"{selected_res.author} {selected_res.title}")
         )
     elif selected_res.index_id in SAUCENAO_DB["anime"]:  # type: ignore
-        final_res.extend(await whatanime_search(url, client))
+        return final_res, whatanime_search
 
-    return final_res
+    return final_res, None
 
 
 async def handle_saucenao_low_acc(
-    url: str,
-    client: ClientSession,
-    mode: str,
-    selected_res: SauceNAOItem,
-) -> List[str]:
-    final_res = []
+    mode: str, selected_res: SauceNAOItem
+) -> Tuple[List[str], Optional[SEARCH_FUNCTION_TYPE]]:
+    final_res: List[str] = []
     # 因为 saucenao 的动画搜索数据库更新不够快，所以当搜索模式为动画时额外增加 whatanime 的搜索结果
     if mode == "anime":
-        final_res.extend(await whatanime_search(url, client))
+        return final_res, whatanime_search
     elif config.auto_use_ascii2d:
         final_res.append(f"相似度 {selected_res.similarity}% 过低，自动使用 Ascii2D 进行搜索")
-        final_res.extend(await ascii2d_search(url, client))
+        return final_res, ascii2d_search
 
-    return final_res
+    return final_res, None