✨ 添加简单统计功能

nonebot_bison/config/db_config.py

@@ -3,6 +3,7 @@
 from datetime import time, datetime
 from collections.abc import Callable, Sequence, Awaitable
 
+from loguru import logger
 from nonebot.compat import model_dump
 from sqlalchemy.orm import selectinload
 from sqlalchemy.exc import IntegrityError
@@ -32,9 +33,11 @@ def __init__(self):
         self.delete_target_hook: list[Callable[[str, T_Target], Awaitable]] = []
 
     def register_add_target_hook(self, fun: Callable[[str, T_Target], Awaitable]):
+        logger.debug(f"register add target hook {fun.__name__}")
         self.add_target_hook.append(fun)
 
     def register_delete_target_hook(self, fun: Callable[[str, T_Target], Awaitable]):
+        logger.debug(f"register delete target hook {fun.__name__}")
         self.delete_target_hook.append(fun)
 
     async def add_subscribe(
@@ -55,8 +58,10 @@ async def add_subscribe(
             db_target_stmt = select(Target).where(Target.platform_name == platform_name).where(Target.target == target)
             db_target: Target | None = await session.scalar(db_target_stmt)
             if not db_target:
+                logger.debug(f"add sub get db_target: {db_target}")
                 db_target = Target(target=target, platform_name=platform_name, target_name=target_name)
-                await asyncio.gather(*[hook(platform_name, target) for hook in self.add_target_hook])
+                hook_resp = await asyncio.gather(*[hook(platform_name, target) for hook in self.add_target_hook])
+                logger.debug(f"add hook response: {hook_resp}")
             else:
                 db_target.target_name = target_name
             subscribe = Subscribe(
@@ -106,7 +111,10 @@ async def del_subscribe(self, user: PlatformTarget, target: str, platform_name:
             )
             if target_count == 0:
                 # delete empty target
-                await asyncio.gather(*[hook(platform_name, T_Target(target)) for hook in self.delete_target_hook])
+                hook_resp = await asyncio.gather(
+                    *[hook(platform_name, T_Target(target)) for hook in self.delete_target_hook]
+                )
+                logger.debug(f"del hook response: {hook_resp}")
             await session.commit()
 
     async def update_subscribe(

nonebot_bison/scheduler/manager.py

@@ -1,16 +1,21 @@
 from typing import cast
 
+from nonebot import on_command
 from nonebot.log import logger
+from nonebot.rule import to_me
+from nonebot.permission import SUPERUSER
 
 from ..utils import Site
 from ..config import config
 from .scheduler import Scheduler
 from ..config.db_model import Target
 from ..types import Target as T_Target
 from ..platform import platform_manager
+from .statistic import runtime_statistic
 from ..plugin_config import plugin_config
 from ..utils.site import CookieClientManager, is_cookie_client_manager
 
+inspect_scheduler = on_command("inspect-bison", priority=5, rule=to_me(), permission=SUPERUSER)
 scheduler_dict: dict[type[Site], Scheduler] = {}
 
 
@@ -63,3 +68,10 @@ async def handle_delete_target(platform_name: str, target: T_Target):
     platform = platform_manager[platform_name]
     scheduler_obj = scheduler_dict[platform.site]
     scheduler_obj.delete_schedulable(platform_name, target)
+
+
+@inspect_scheduler.handle()
+async def inspect_scheduler_handle():
+    report = runtime_statistic.generate_report(scheduler_dict)
+    await inspect_scheduler.send("统计数据仅为本次启动后的数据，不包含历史数据")
+    await inspect_scheduler.finish(report)

nonebot_bison/scheduler/scheduler.py

@@ -12,6 +12,7 @@
 from ..types import Target, SubUnit
 from ..platform import platform_manager
 from ..utils import Site, ProcessContext
+from .statistic import runtime_statistic
 from ..utils.site import SkipRequestException
 
 
@@ -71,6 +72,7 @@ def _refresh_batch_api_target_cache(self):
             for target in targets:
                 self.batch_api_target_cache[platform_name][target] = targets
 
+    @runtime_statistic.statistic_schedule_count
     async def get_next_schedulable(self) -> Schedulable | None:
         if not self.schedulable_list:
             return None
@@ -124,13 +126,15 @@ async def exec_fetch(self):
             for send_post in send_list:
                 logger.info(f"send to {user}: {send_post}")
                 try:
+                    runtime_statistic.statistic_post_send(send_post.platform.name)
                     await send_msgs(
                         user,
                         await send_post.generate_messages(),
                     )
                 except NoBotFound:
                     logger.warning("no bot connected")
 
+    @runtime_statistic.statistic_record("insert_new")
     def insert_new_schedulable(self, platform_name: str, target: Target):
         self.pre_weight_val += 1000
         new_schedulable = Schedulable(platform_name, target, 1000, platform_manager[platform_name].use_batch)
@@ -142,6 +146,7 @@ def insert_new_schedulable(self, platform_name: str, target: Target):
         self.schedulable_list.append(new_schedulable)
         logger.info(f"insert [{platform_name}]{target} to Schduler({self.scheduler_config.name})")
 
+    @runtime_statistic.statistic_record("delete")
     def delete_schedulable(self, platform_name, target: Target):
         if platform_manager[platform_name].use_batch:
             self.batch_platform_name_targets_cache[platform_name].remove(target)
@@ -157,3 +162,5 @@ def delete_schedulable(self, platform_name, target: Target):
         if to_find_idx is not None:
             deleted_schdulable = self.schedulable_list.pop(to_find_idx)
             self.pre_weight_val -= deleted_schdulable.current_weight
+
+        logger.info(f"delete [{platform_name}]{target} from Schduler({self.scheduler_config.name})")

nonebot_bison/scheduler/statistic.py

@@ -0,0 +1,99 @@
+from functools import wraps
+from datetime import datetime
+from collections import defaultdict
+from collections.abc import Callable, Coroutine
+from typing import TYPE_CHECKING, Any, Literal, TypeVar, TypedDict
+
+from nonebot_bison.utils import dict_to_str
+
+from ..utils import Site
+from ..types import Target
+
+if TYPE_CHECKING:
+    from .scheduler import Scheduler, Schedulable
+
+
+TSchd = TypeVar("TSchd", bound="Scheduler")
+
+
+RecordKind = Literal["insert_new", "delete"]
+
+
+class PlatformTargetRecord(TypedDict):
+    schedule_count: defaultdict[str, defaultdict[str, int]]
+    insert_new: list[tuple[str, datetime]]
+    delete: list[tuple[str, datetime]]
+    post_send: defaultdict[str, int]
+
+
+class RuntimeStatistic:
+    def __init__(self):
+        self._record = PlatformTargetRecord(
+            schedule_count=defaultdict(lambda: defaultdict(int)), insert_new=[], delete=[], post_send=defaultdict(int)
+        )
+
+    def statistic_schedule_count(self, func: "Callable[[TSchd], Coroutine[Any, Any, Schedulable | None]]"):
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            if not (schedulable := await func(*args, **kwargs)):
+                return
+            self._record["schedule_count"][schedulable.platform_name][schedulable.target] += 1
+            return schedulable
+
+        return wrapper
+
+    def statistic_record(self, name: RecordKind):
+        record = self._record[name]
+
+        def decorator(func: "Callable[[TSchd, str, Target], None]"):
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                platform_name, target = args[1], args[2]
+                record.append((f"{platform_name}-{target}", datetime.now()))
+                return func(*args, **kwargs)
+
+            return wrapper
+
+        return decorator
+
+    def statistic_post_send(self, platform_name: str):
+        self._record["post_send"][platform_name] += 1
+
+    def _generate_stats(self, scheduler_dict: "dict[type[Site], Scheduler]"):
+        report_dict: dict[str, Any] = {
+            "新增订阅hook调用记录": [f"{record[0]}: {record[1]}" for record in self._record["insert_new"]],
+            "删除订阅hook调用记录": [f"{record[0]}: {record[1]}" for record in self._record["delete"]],
+        }
+
+        platform_dict = {}
+        for platform, targets in self._record["schedule_count"].items():
+            target_dict = [f"{target}: {count} 次" for target, count in targets.items()]
+            platform_dict[platform] = target_dict
+        report_dict["调度统计"] = platform_dict
+
+        post_send_dict = {}
+        for platform, count in self._record["post_send"].items():
+            post_send_dict[platform] = f"{count} 次"
+
+        report_dict["发送消息统计"] = post_send_dict
+
+        all_schedulable = {}
+        for site, scheduler in scheduler_dict.items():
+            scheduler_list = []
+            for schedulable in scheduler.schedulable_list:
+                scheduler_list.append(
+                    f"{schedulable.target}: [权重]{schedulable.current_weight} | [批量]{schedulable.use_batch}"
+                )
+            all_schedulable[site.name] = scheduler_list
+
+        report_dict["所有调度对象"] = all_schedulable
+
+        return report_dict
+
+    def generate_report(self, scheduler_dict: "dict[type[Site], Scheduler]"):
+        """根据 schduler_dict 生成报告"""
+        report_dict = self._generate_stats(scheduler_dict)
+        return dict_to_str(report_dict)
+
+
+runtime_statistic = RuntimeStatistic()

nonebot_bison/sub_manager/del_sub.py

@@ -1,3 +1,4 @@
+from loguru import logger
 from nonebot.typing import T_State
 from nonebot.matcher import Matcher
 from nonebot.params import Arg, EventPlainText
@@ -50,6 +51,7 @@ async def do_del(
             index = int(index_str)
             await config.del_subscribe(user_info, **state["sub_table"][index])
         except Exception:
+            logger.exception("删除订阅错误")
             await del_sub.reject("删除错误")
         else:
             await del_sub.finish("删除成功")

nonebot_bison/utils/__init__.py

@@ -1,6 +1,7 @@
 import re
 import sys
 import difflib
+from typing import Any
 
 import nonebot
 from nonebot.plugin import require
@@ -114,3 +115,36 @@ def decode_match(match: re.Match[str]) -> str:
 def text_fletten(text: str, *, banned: str = "\n\r\t", replace: str = " ") -> str:
     """将文本中的格式化字符去除"""
     return "".join(c if c not in banned else replace for c in text)
+
+
+def dict_to_str(__dict: dict[str, Any]) -> str:
+    """根据 dict 生成带缩进的文本，
+    缩进层级由 dict 的嵌套关系决定，
+    每层的 key 为标题，value 为内容，
+    内容为列表时每个元素为一行，内容为字典时递归生成，内容为其他类型时直接输出，
+    生成的报告字符串应该是一个完整的文档，包含标题和内容
+    每层缩进 2 个空格
+    例如：{"a": ["b", "c"], "d": {"e": ["f", "g"]}} 生成的报告为：
+    a:
+        b
+        c
+    d:
+        e:
+        f
+        g
+    """
+
+    def generate_report(_d, level=0):
+        res = ""
+        for key, value in _d.items():
+            res += "  " * level + key + ":\n"
+            if isinstance(value, list):
+                for item in value:
+                    res += "  " * (level + 1) + item + "\n"
+            elif isinstance(value, dict):
+                res += generate_report(value, level + 1)
+            else:
+                res += "  " * (level + 1) + str(value) + "\n"
+        return res
+
+    return generate_report(__dict)

tests/scheduler/test_statistic.py

@@ -0,0 +1,112 @@
+from textwrap import dedent
+from datetime import datetime
+from typing import TYPE_CHECKING
+from collections import defaultdict
+from unittest.mock import Mock, AsyncMock, patch
+
+import pytest
+
+if TYPE_CHECKING:
+    from nonebot_bison.scheduler.statistic import RuntimeStatistic
+
+
+@pytest.fixture
+def runtime_statistic():
+    from nonebot_bison.scheduler.statistic import RuntimeStatistic
+
+    return RuntimeStatistic()
+
+
+@pytest.mark.asyncio
+async def test_statistic_schedule_count(app, runtime_statistic: "RuntimeStatistic"):
+    mock_func = AsyncMock()
+    mock_func.return_value = AsyncMock(platform_name="test_platform", target="test_target")
+
+    decorated_func = runtime_statistic.statistic_schedule_count(mock_func)
+    await decorated_func()
+
+    assert runtime_statistic._record["schedule_count"]["test_platform"]["test_target"] == 1
+
+
+@pytest.mark.asyncio
+async def test_statistic_schedule_count_no_schedulable(app, runtime_statistic: "RuntimeStatistic"):
+    mock_func = AsyncMock()
+    mock_func.return_value = None
+
+    decorated_func = runtime_statistic.statistic_schedule_count(mock_func)
+    await decorated_func()
+
+    assert "test_platform" not in runtime_statistic._record["schedule_count"]
+
+
+@pytest.mark.asyncio
+async def test_statistic_record_insert_new(app, runtime_statistic: "RuntimeStatistic"):
+    mock_func = AsyncMock()
+    mock_self = Mock()
+
+    decorated_func = runtime_statistic.statistic_record("insert_new")(mock_func)
+    with patch("nonebot_bison.scheduler.statistic.datetime") as mock_datetime:
+        mock_datetime.now.return_value = datetime(2023, 1, 1)
+        await decorated_func(mock_self, "test_platform", "test_target")  # type: ignore
+
+    assert runtime_statistic._record["insert_new"] == [("test_platform-test_target", datetime(2023, 1, 1))]
+
+
+@pytest.mark.asyncio
+async def test_statistic_record_delete(app, runtime_statistic: "RuntimeStatistic"):
+    mock_func = AsyncMock()
+    mock_self = Mock()
+
+    decorated_func = runtime_statistic.statistic_record("delete")(mock_func)
+    with patch("nonebot_bison.scheduler.statistic.datetime") as mock_datetime:
+        mock_datetime.now.return_value = datetime(2023, 1, 2)
+        await decorated_func(mock_self, "test_platform", "test_target")  # type: ignore
+
+    assert runtime_statistic._record["delete"] == [("test_platform-test_target", datetime(2023, 1, 2))]
+
+
+@pytest.mark.asyncio
+async def test_statistic_record_report_generate(app, runtime_statistic: "RuntimeStatistic"):
+    from nonebot_bison.utils import Site
+    from nonebot_bison.scheduler.scheduler import Scheduler
+
+    class MockSite(Site):
+        name = "test_site"
+        schedule_type = "interval"
+        schedule_setting = {"seconds": 100}
+
+    scheduler_dict: dict[type[Site], Scheduler] = {MockSite: Scheduler(MockSite, [], [])}
+    _mock_record = defaultdict(int)
+    _mock_record["test_target"] = 1
+    runtime_statistic._record["schedule_count"]["test_platform"] = _mock_record
+    runtime_statistic._record["insert_new"] = [("test_platform-test_target", datetime(2023, 1, 1))]
+    runtime_statistic._record["delete"] = [("test_platform-test_target", datetime(2023, 1, 1))]
+
+    report = runtime_statistic._generate_stats(scheduler_dict)
+
+    assert report == {
+        "新增订阅hook调用记录": [f"test_platform-test_target: {datetime(2023, 1, 1).strftime('%Y-%m-%d %H:%M:%S')}"],
+        "删除订阅hook调用记录": [f"test_platform-test_target: {datetime(2023, 1, 1).strftime('%Y-%m-%d %H:%M:%S')}"],
+        "调度统计": {"test_platform": ["test_target: 1 次"]},
+        "所有调度对象": {"test_site": []},
+        "发送消息统计": {},
+    }
+
+    repost_str = runtime_statistic.generate_report(scheduler_dict)
+    assert (
+        repost_str
+        == dedent(
+            """
+        新增订阅hook调用记录:
+          test_platform-test_target: 2023-01-01 00:00:00
+        删除订阅hook调用记录:
+          test_platform-test_target: 2023-01-01 00:00:00
+        调度统计:
+          test_platform:
+            test_target: 1 次
+        发送消息统计:
+        所有调度对象:
+          test_site:
+        """
+        )[1:]
+    )  # 不要第一个换行符