♻️ 重构cd系统 (#231)

Well2333 · Aug 4, 2024 · 4488aad · 4488aad
1 parent ebb8ed8
commit 4488aad
Show file tree

Hide file tree

Showing 8 changed files with 69 additions and 42 deletions.
diff --git a/bot.py b/bot.py
@@ -22,6 +22,13 @@
     compression="tar.xz",
     retention="7 days",
 )
+logger.add(
+    "logs/trace/{time:YYYY-MM-DD}.log",
+    level="TRACE",
+    rotation="00:00",
+    compression="tar.xz",
+    retention="7 days",
+)
 logger.info("========= 重新启动 =========")
 
 driver = nonebot.get_driver()

diff --git a/nonebot_plugin_bilichat/base_content_parsing.py b/nonebot_plugin_bilichat/base_content_parsing.py
@@ -1,7 +1,6 @@
 import asyncio
 import re
 import shlex
-import time
 
 from nonebot.adapters import Bot, Event
 from nonebot.exception import FinishedException
@@ -14,6 +13,7 @@
 from .config import plugin_config
 from .content import Column, Dynamic, Video
 from .lib.b23_extract import b23_extract
+from .lib.content_cd import BilichatCD
 from .lib.text_to_image import t2i
 from .model.arguments import Options, parser
 from .model.exception import AbortError, ProssesError
@@ -25,34 +25,9 @@
 if plugin_config.bilichat_word_cloud:
     from .wordcloud import wordcloud
 
-cd: dict[str, int] = {}
-cd_size_limit = plugin_config.bilichat_cd_time // 2
 lock = asyncio.Lock()
 
 
-def check_cd(uid: int | str, check: bool = True):
-    global cd
-    now = int(time.time())
-    uid = str(uid)
-    # gc
-    if len(cd) > cd_size_limit:
-        cd = {k: cd[k] for k in cd if cd[k] < now}
-    # not check cd
-    if not check:
-        cd[uid] = now + plugin_config.bilichat_cd_time
-        return
-    # check cd
-    session, id_ = uid.split("_-_")
-    if cd.get(uid, 0) > now:
-        logger.warning(f"会话 [{session}] 的重复内容 [{id_}]. 跳过解析")
-        raise FinishedException
-    elif cd.get(id_, 0) > now:
-        logger.warning(f"会话 [全局] 的重复内容 [{id_}]. 跳过解析")
-        raise FinishedException
-    else:
-        cd[uid] = now + plugin_config.bilichat_cd_time
-
-
 async def _permission_check(bot: Bot, event: Event, target: MsgTarget, state: T_State):
     state["_uid_"] = target.id
     # 自身消息
@@ -124,7 +99,10 @@ async def _bili_check(state: T_State, event: Event, bot: Bot, msg: UniMsg) -> bo
             content = await Dynamic.from_id(_id)
 
         if content:
-            check_cd(f"{state['_uid_']}_-_{content.id}", check=not options.force)
+            if options.force:
+                BilichatCD.record_cd(state["_uid_"], str(content.id))
+            else:
+                BilichatCD.check_cd(state["_uid_"], str(content.id))
             state["_content_"] = content
         else:
             raise AbortError(f"查询 {bililink} 返回内容为空")

diff --git a/nonebot_plugin_bilichat/commands/functions.py b/nonebot_plugin_bilichat/commands/functions.py
@@ -7,10 +7,10 @@
 from nonebot.typing import T_State
 from nonebot_plugin_alconna.uniseg import Hyper, Image, MsgTarget, Reply, Text, UniMessage, UniMsg
 
-from ..base_content_parsing import check_cd
 from ..config import plugin_config
 from ..content import Column, Dynamic, Video
 from ..lib.b23_extract import b23_extract
+from ..lib.content_cd import BilichatCD
 from ..lib.fetch_dynamic import fetch_last_dynamic
 from ..lib.uid_extract import uid_extract
 from ..model.exception import AbortError
@@ -28,6 +28,7 @@ async def check_dynamic_v11(target: MsgTarget, uid: Message = CommandArg()):
     if isinstance(up, str):
         await bili_check_dyn.finish(up)
     if dyn := await fetch_last_dynamic(up):
+        BilichatCD.record_cd(session_id=target.id, content_id=dyn.id)
         if image := await dyn.get_image(plugin_config.bilichat_dynamic_style):
             await UniMessage(Image(raw=image)).send(target=target)
 
@@ -73,7 +74,7 @@ async def fetch_check(state: T_State, msg: UniMsg, target: MsgTarget):
             raise AbortError("该功能目前仅可用于图文动态哦~")
 
         if content:
-            check_cd(f"{target.id}_-_{content.id}", check=False)
+            BilichatCD.record_cd(session_id=target.id, content_id=content.id)
             state["_content_"] = content
         else:
             raise AbortError(f"查询 {bililink} 返回内容为空")

diff --git a/nonebot_plugin_bilichat/lib/content_cd.py b/nonebot_plugin_bilichat/lib/content_cd.py
@@ -0,0 +1,48 @@
+from datetime import datetime, timedelta
+
+from nonebot.exception import FinishedException
+from nonebot.log import logger
+
+from ..config import plugin_config
+
+
+class BilichatCD:
+    cd: dict[str, dict[str, datetime]] = {}  # {content_id: {session_id: datetime_to_expire}}
+    cd_size_limit = plugin_config.bilichat_cd_time // 2
+    expiration_duration = timedelta(seconds=plugin_config.bilichat_cd_time)
+
+    @classmethod
+    def check_cd(cls, session_id: str, content_id: str):
+        logger.trace(f"当前记录信息: {cls.cd}")
+        logger.trace(f"content:{content_id} session:{session_id}")
+        content_record = cls.cd.get(content_id, {})
+        now = datetime.now()
+
+        if session_id in content_record and content_record[session_id] > now:
+            logger.warning(f"会话 [{session_id}] 的重复内容 [{content_id}]. 跳过解析")
+            raise FinishedException
+        elif "global" in content_record and content_record["global"] > now:
+            logger.warning(f"会话 [全局] 的重复内容 [{content_id}]. 跳过解析")
+            raise FinishedException
+        else:
+            cls.record_cd(session_id, content_id)
+
+    @classmethod
+    def record_cd(cls, session_id: str, content_id: str):
+        content_record = cls.cd.get(content_id, {})
+        now = datetime.now()
+
+        # Clean up expired entries
+        cls.clean_expired_entries(content_record, now)
+
+        # Record new entry
+        content_record[session_id] = now + cls.expiration_duration
+        cls.cd[content_id] = content_record
+
+    @classmethod
+    def clean_expired_entries(cls, content_record: dict[str, datetime], current_time: datetime):
+        expired_entries = [
+            session_id for session_id, expire_time in content_record.items() if expire_time <= current_time
+        ]
+        for session_id in expired_entries:
+            del content_record[session_id]
diff --git a/nonebot_plugin_bilichat/lib/fetch_dynamic.py b/nonebot_plugin_bilichat/lib/fetch_dynamic.py
@@ -7,7 +7,6 @@
 from httpx import TimeoutException
 from nonebot.log import logger
 
-from ..base_content_parsing import check_cd
 from ..content.dynamic import Dynamic
 from ..lib.bilibili_request import get_b23_url, get_user_dynamics
 from ..lib.bilibili_request.auth import AuthManager
@@ -48,7 +47,6 @@ async def _fetch_rest(up_mid: int, up_name: str) -> Dynamic | None:
     dyns = {int(x["id_str"]): x for x in resp if x["type"] not in DYNAMIC_TYPE_IGNORE}
     dyn_id = max(dyns.keys())
     dyn = dyns[dyn_id]
-    check_cd(dyn["id_str"], check=False)
 
     url = await get_b23_url(f"https://t.bilibili.com/{dyn['id_str']}")
     dynamic = Dynamic(id=dyn["id_str"], url=url, raw=dyn, raw_type="web")
@@ -77,7 +75,6 @@ async def _fetch_grpc(up_mid: int, up_name: str) -> Dynamic | None:
     dyns = {int(x.extend.dyn_id_str): x for x in resp.list if x.card_type not in DYNAMIC_TYPE_IGNORE}
     dyn_id = max(dyns.keys())
     dyn = dyns[dyn_id]
-    check_cd(dyn.extend.dyn_id_str, check=False)
 
     url = await get_b23_url(f"https://t.bilibili.com/{dyn.extend.dyn_id_str}")
     dynamic = Dynamic(

diff --git a/nonebot_plugin_bilichat/subscribe/dynamic.py b/nonebot_plugin_bilichat/subscribe/dynamic.py
@@ -11,12 +11,12 @@
 from httpx import TimeoutException
 from nonebot.log import logger
 
-from ..base_content_parsing import check_cd
 from ..config import plugin_config
 from ..content.dynamic import Dynamic
 from ..lib.bilibili_request import get_b23_url, get_user_dynamics
 from ..lib.bilibili_request.auth import AuthManager
-from ..model.const import DYNAMIC_TYPE_MAP, DYNAMIC_TYPE_IGNORE
+from ..lib.content_cd import BilichatCD
+from ..model.const import DYNAMIC_TYPE_IGNORE, DYNAMIC_TYPE_MAP
 from ..model.exception import AbortError
 from ..optional import capture_exception
 from .manager import Uploader
@@ -45,7 +45,6 @@ async def fetch_dynamics_rest(up: Uploader):
     dyns = [x for x in resp if int(x["id_str"]) > up.dyn_offset and x["type"] not in DYNAMIC_TYPE_IGNORE]
     dyns.reverse()
     for dyn in dyns:
-        check_cd(dyn["id_str"], check=False)
         up.dyn_offset = max(up.dyn_offset, int(dyn["id_str"]))
         up_name = dyn["modules"]["module_author"]["name"]
         if up.nickname != up_name:
@@ -59,12 +58,10 @@ async def fetch_dynamics_rest(up: Uploader):
             type_text += "投稿了视频"
             aid = dyn["modules"]["module_dynamic"]["major"]["archive"]["aid"]
             url = await get_b23_url(f"https://www.bilibili.com/video/av{aid}")
-            check_cd(aid, check=False)
         elif dyn_type == DynamicType.article:
             type_text += "投稿了专栏"
             cvid = dyn["modules"]["module_dynamic"]["major"]["article"]["id"]
             url = await get_b23_url(f"https://www.bilibili.com/read/cv{cvid}")
-            check_cd(cvid, check=False)
         elif dyn_type == DynamicType.music:
             type_text += "投稿了音乐"
         elif dyn_type == DynamicType.forward:
@@ -87,6 +84,7 @@ async def fetch_dynamics_rest(up: Uploader):
         content = [type_text, dyn_image, url]
         for user in up.subscribed_users:
             if user.subscriptions[up.uid].dynamic:
+                BilichatCD.record_cd(session_id=user.user_id, content_id=dynamic.id)
                 await user.push_to_user(
                     content=content, at_all=user.subscriptions[up.uid].dynamic_at_all or user.at_all
                 )
@@ -117,7 +115,6 @@ async def fetch_dynamics_grpc(up: Uploader):
     dyns = [x for x in resp.list if int(x.extend.dyn_id_str) > up.dyn_offset and x.card_type not in DYNAMIC_TYPE_IGNORE]
     dyns.reverse()
     for dyn in dyns:
-        check_cd(dyn.extend.dyn_id_str, check=False)
         up.dyn_offset = max(up.dyn_offset, int(dyn.extend.dyn_id_str))
         up_name = dyn.modules[0].module_author.author.name
         if up.nickname != up_name:
@@ -132,14 +129,12 @@ async def fetch_dynamics_grpc(up: Uploader):
                 if module.module_type == DynModuleType.module_dynamic:
                     aid = module.module_dynamic.dyn_archive.avid
                     url = await get_b23_url(f"https://www.bilibili.com/video/av{aid}")
-                    check_cd(aid, check=False)
         elif dyn.card_type == DynamicType.article:
             type_text += "投稿了专栏"
             for module in dyn.modules:
                 if module.module_type == DynModuleType.module_dynamic:
                     cvid = module.module_dynamic.dyn_article.id
                     url = await get_b23_url(f"https://www.bilibili.com/read/cv{cvid}")
-                    check_cd(cvid, check=False)
         elif dyn.card_type == DynamicType.music:
             type_text += "投稿了音乐"
         elif dyn.card_type == DynamicType.forward:
@@ -165,6 +160,7 @@ async def fetch_dynamics_grpc(up: Uploader):
         content = [type_text, dyn_image, url]
         for user in up.subscribed_users:
             if user.subscriptions[up.uid].dynamic:
+                BilichatCD.record_cd(session_id=user.user_id, content_id=dynamic.id)
                 await user.push_to_user(
                     content=content, at_all=user.subscriptions[up.uid].dynamic_at_all or user.at_all
                 )
diff --git a/pdm.lock b/pdm.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "nonebot-plugin-bilichat"
 
-version = "5.10.4"
+version = "5.11.0"
 
 description = "多种B站链接解析，视频词云，AI总结，你想要的都在 bilichat"
 authors = [