bugfix:修复notice事件扩散问题 (#2132)

* bugfix:修复notice事件扩散问题

* 优化并发调度

* bugfix:修复签到样式

* bugfix：功能调用统计修复

* bugfix:修复私聊时功能调用统计显示已退群问题

* 提高插件适配兼容性

* 优化发送队列

* 修改权限检查设计

* 继续修改权限检查设计

* 完善权限检查设计

* 优化sqlite配置

* 优化数据库初始化

* 代码整理，无用代码清理

* bugfix:修复启动时数据库校验问题

* bugfix:修复预算裁剪过于激进问题

zhenxun/builtin_plugins/hooks/auth/auth_limit.py

@@ -1,4 +1,6 @@
 import asyncio
+from collections.abc import Callable
+from dataclasses import dataclass, field
 import time
 from typing import Any, ClassVar
 
@@ -46,6 +48,26 @@ class Limit(BaseModel):
         arbitrary_types_allowed = True
 
 
+@dataclass(slots=True)
+class LimitReservation:
+    module: str
+    releases: list[Callable[[], None]] = field(default_factory=list)
+    should_auto_unblock: bool = False
+    active: bool = True
+
+    def commit(self) -> None:
+        self.active = False
+        self.releases.clear()
+
+    def release(self) -> None:
+        if not self.active:
+            return
+        for release in reversed(self.releases):
+            release()
+        self.active = False
+        self.releases.clear()
+
+
 def _limit_notice_key(
     limit: PluginLimit | PluginLimitSnapshot,
     user_id: str,
@@ -247,14 +269,9 @@ class LimitManager:
             for limit in limits:
                 cls.add_limit(limit)
 
-        # 检查各种限制
         try:
-            if limit_model := cls.cd_limit.get(module):
-                await cls.__check(limit_model, user_id, group_id, channel_id)
-            if limit_model := cls.block_limit.get(module):
-                await cls.__check(limit_model, user_id, group_id, channel_id)
-            if limit_model := cls.count_limit.get(module):
-                await cls.__check(limit_model, user_id, group_id, channel_id)
+            reservation = await cls.reserve(module, user_id, group_id, channel_id)
+            reservation.commit()
         finally:
             # 记录总执行时间
             elapsed = time.time() - start_time
@@ -267,13 +284,53 @@ class LimitManager:
                 )
 
     @classmethod
-    async def __check(
+    async def reserve(
+        cls,
+        module: str,
+        user_id: str,
+        group_id: str | None,
+        channel_id: str | None,
+    ) -> LimitReservation:
+        """检查并预留限制状态；调用方失败时可 release 回滚内存限制。"""
+        if (
+            time.time() - cls.last_update_time > cls.update_interval
+            and not cls.is_updating
+        ):
+            asyncio.create_task(cls.update_limits())  # noqa: RUF006
+
+        if module not in cls.add_module:
+            limits = await cls.get_module_limits(module)
+            for limit in limits:
+                cls.add_limit(limit)
+
+        reservation = LimitReservation(module=module)
+        try:
+            if limit_model := cls.cd_limit.get(module):
+                reservation.releases.append(
+                    await cls.__reserve(limit_model, user_id, group_id, channel_id)
+                )
+            if limit_model := cls.block_limit.get(module):
+                reservation.should_auto_unblock = True
+                reservation.releases.append(
+                    await cls.__reserve(limit_model, user_id, group_id, channel_id)
+                )
+            if limit_model := cls.count_limit.get(module):
+                reservation.releases.append(
+                    await cls.__reserve(limit_model, user_id, group_id, channel_id)
+                )
+        except Exception:
+            reservation.release()
+            raise
+        return reservation
+
+    @classmethod
+    async def __reserve(
         cls,
         limit_model: Limit | None,
         user_id: str,
         group_id: str | None,
         channel_id: str | None,
-    ):
+    ) -> Callable[[], None]:
         """检测限制
 
         参数:
@@ -286,7 +343,7 @@ class LimitManager:
             IgnoredException: IgnoredException
         """
         if not limit_model:
-            return
+            return lambda: None
         limit = limit_model.limit
         limiter = limit_model.limiter
         is_limit = (
@@ -317,12 +374,40 @@ class LimitManager:
                 group_id=group_id,
             )
             if isinstance(limiter, FreqLimiter):
+                had_next_time = key_type in limiter.next_time
+                old_next_time = limiter.next_time.get(key_type, 0.0)
                 limiter.start_cd(key_type)
+
+                def release_freq() -> None:
+                    if had_next_time:
+                        limiter.next_time[key_type] = old_next_time
+                    else:
+                        limiter.next_time.pop(key_type, None)
+
+                return release_freq
             if isinstance(limiter, UserBlockLimiter):
+                old_flag = limiter.flag_data.get(key_type, False)
+                old_time = limiter.time.get(key_type, 0.0)
                 limiter.set_true(key_type)
+
+                def release_block() -> None:
+                    limiter.flag_data[key_type] = old_flag
+                    if old_time:
+                        limiter.time[key_type] = old_time
+                    else:
+                        limiter.time.pop(key_type, None)
+
+                return release_block
             if isinstance(limiter, CountLimiter):
+                old_count = limiter.count.get(key_type, 0)
                 limiter.increase(key_type)
 
+                def release_count() -> None:
+                    limiter.count[key_type] = old_count
+
+                return release_count
+        return lambda: None
+
 
 async def auth_limit(
     plugin: PluginInfo,
@@ -343,11 +428,42 @@ async def auth_limit(
         entity = get_entity_ids(session)
     try:
         await asyncio.wait_for(
-            LimitManager.check(
-                plugin.module, entity.user_id, entity.group_id, entity.channel_id
-            ),
+            _reserve_and_commit_limit(plugin.module, entity),
             timeout=DB_TIMEOUT_SECONDS * 2,  # 给予更长的超时时间
         )
     except asyncio.TimeoutError:
         logger.error(f"检查插件限制超时: {plugin.module}", LOGGER_COMMAND)
         # 超时时不抛出异常，允许继续执行
+
+
+async def reserve_auth_limit(
+    plugin: PluginInfo,
+    session: Uninfo,
+    *,
+    context: PermissionContext | None = None,
+    entity: EntityIDs | None = None,
+) -> LimitReservation:
+    del session
+    if context is not None:
+        entity = context.entity
+    if entity is None:
+        raise RuntimeError("reserve_auth_limit requires entity or context")
+    return await LimitManager.reserve(
+        plugin.module,
+        entity.user_id,
+        entity.group_id,
+        entity.channel_id,
+    )
+
+
+async def _reserve_and_commit_limit(
+    module: str,
+    entity: EntityIDs,
+) -> None:
+    reservation = await LimitManager.reserve(
+        module,
+        entity.user_id,
+        entity.group_id,
+        entity.channel_id,
+    )
+    reservation.commit()