feat: 恢复ratelimit

pkg/pipeline/ratelimit/algo.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+import abc
+
+from ...core import app
+
+
+class ReteLimitAlgo(metaclass=abc.ABCMeta):
+    
+    ap: app.Application
+
+    def __init__(self, ap: app.Application):
+        self.ap = ap
+
+    async def initialize(self):
+        pass
+
+    @abc.abstractmethod
+    async def require_access(self, launcher_type: str, launcher_id: int) -> bool:
+        raise NotImplementedError
+    
+    @abc.abstractmethod
+    async def release_access(self, launcher_type: str, launcher_id: int):
+        raise NotImplementedError
+ 

pkg/pipeline/ratelimit/algos/fixedwin.py

@@ -0,0 +1,85 @@
+# 固定窗口算法
+from __future__ import annotations
+
+import asyncio
+import time
+
+from .. import algo
+
+
+class SessionContainer:
+    
+    wait_lock: asyncio.Lock
+
+    records: dict[int, int]
+    """访问记录，key为每分钟的起始时间戳，value为访问次数"""
+
+    def __init__(self):
+        self.wait_lock = asyncio.Lock()
+        self.records = {}
+
+
+class FixedWindowAlgo(algo.ReteLimitAlgo):
+
+    containers_lock: asyncio.Lock
+    """访问记录容器锁"""
+
+    containers: dict[str, SessionContainer]
+    """访问记录容器，key为launcher_type launcher_id"""
+
+    async def initialize(self):
+        self.containers_lock = asyncio.Lock()
+        self.containers = {}
+
+    async def require_access(self, launcher_type: str, launcher_id: int) -> bool:
+        # 加锁，找容器
+        container: SessionContainer = None
+
+        session_name = f'{launcher_type}_{launcher_id}'
+
+        async with self.containers_lock:
+            container = self.containers.get(session_name)
+
+            if container is None:
+                container = SessionContainer()
+                self.containers[session_name] = container
+
+        # 等待锁
+        async with container.wait_lock:
+            # 获取当前时间戳
+            now = int(time.time())
+
+            # 获取当前分钟的起始时间戳
+            now = now - now % 60
+
+            # 获取当前分钟的访问次数
+            count = container.records.get(now, 0)
+
+            limitation = self.ap.cfg_mgr.data['rate_limitation']['default']
+
+            if session_name in self.ap.cfg_mgr.data['rate_limitation']:
+                limitation = self.ap.cfg_mgr.data['rate_limitation'][session_name]
+
+            # 如果访问次数超过了限制
+            if count >= limitation:
+                if self.ap.cfg_mgr.data['rate_limit_strategy'] == 'drop':
+                    return False
+                elif self.ap.cfg_mgr.data['rate_limit_strategy'] == 'wait':
+                    # 等待下一分钟
+                    await asyncio.sleep(60 - time.time() % 60)
+    
+                    now = int(time.time())
+                    now = now - now % 60
+            
+            if now not in container.records:
+                container.records = {}
+                container.records[now] = 1
+            else:
+                # 访问次数加一
+                container.records[now] = count + 1
+
+            # 返回True
+            return True
+        
+    async def release_access(self, launcher_type: str, launcher_id: int):
+        pass

pkg/pipeline/ratelimit/ratelimit.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+import typing
+
+from .. import entities, stagemgr, stage
+from . import algo
+from .algos import fixedwin
+from ...core import entities as core_entities
+
+
+@stage.stage_class("RequireRateLimitOccupancy")
+@stage.stage_class("ReleaseRateLimitOccupancy")
+class RateLimit(stage.PipelineStage):
+
+    algo: algo.ReteLimitAlgo
+
+    async def initialize(self):
+        self.algo = fixedwin.FixedWindowAlgo(self.ap)
+        await self.algo.initialize()
+
+    async def process(
+        self,
+        query: core_entities.Query,
+        stage_inst_name: str,
+    ) -> typing.Union[
+        entities.StageProcessResult,
+        typing.AsyncGenerator[entities.StageProcessResult, None],
+    ]:
+        """处理
+        """
+        if stage_inst_name == "RequireRateLimitOccupancy":
+            if await self.algo.require_access(
+                query.launcher_type.value,
+                query.launcher_id,
+            ):
+                return entities.StageProcessResult(
+                    result_type=entities.ResultType.CONTINUE,
+                    new_query=query,
+                )
+            else:
+                return entities.StageProcessResult(
+                    result_type=entities.ResultType.INTERRUPT,
+                    new_query=query,
+                    console_notice=f"根据限速规则忽略 {query.launcher_type.value}:{query.launcher_id} 消息",
+                    user_notice=self.ap.tips_mgr.data['rate_limit_drop_tip']
+                )
+        elif stage_inst_name == "ReleaseRateLimitOccupancy":
+            await self.algo.release_access(
+                query.launcher_type,
+                query.launcher_id,
+            )
+            return entities.StageProcessResult(
+                result_type=entities.ResultType.CONTINUE,
+                new_query=query,
+            )

pkg/pipeline/stagemgr.py

@@ -12,6 +12,7 @@ from .longtext import longtext
 from .respback import respback
 from .wrapper import wrapper
 from .preproc import preproc
+from .ratelimit import ratelimit
 
 
 stage_order = [
@@ -19,7 +20,9 @@ stage_order = [
     "BanSessionCheckStage",
     "PreContentFilterStage",
     "PreProcessor",
+    "RequireRateLimitOccupancy",
     "MessageProcessor",
+    "ReleaseRateLimitOccupancy",
     "PostContentFilterStage",
     "ResponseWrapper",
     "LongTextProcessStage",