zhenxun_bot/zhenxun/builtin_plugins/statistics/statistics_hook.py

import asyncio
from datetime import datetime

from nonebot import get_driver
from nonebot.adapters import Bot, Event
from nonebot.adapters.onebot.v11 import PokeNotifyEvent
from nonebot.matcher import Matcher
from nonebot.message import run_postprocessor
from nonebot.plugin import PluginMetadata
from nonebot_plugin_apscheduler import scheduler
from nonebot_plugin_uninfo import Uninfo

from zhenxun.configs.utils import PluginExtraData
from zhenxun.models.statistics import Statistics
from zhenxun.services.cache.runtime_cache import PluginInfoMemoryCache
from zhenxun.services.log import logger
from zhenxun.services.message_load import should_pause_tasks
from zhenxun.utils.enum import PluginType
from zhenxun.utils.utils import get_entity_ids

__plugin_meta__ = PluginMetadata(
    name="功能调用统计",
    description="功能调用统计",
    usage="""""".strip(),
    extra=PluginExtraData(
        author="HibiKier", version="0.1", plugin_type=PluginType.HIDDEN
    ).to_dict(),
)

STATS_BUFFER_FLUSH_SIZE = 5000
STATS_BUFFER_MAX_RETAIN = 10000
TEMP_LIST: list[Statistics] = []
_STATS_FLUSH_LOCK = asyncio.Lock()
driver = get_driver()


async def _flush_statistics_buffer(reason: str) -> int:
    async with _STATS_FLUSH_LOCK:
        call_list = TEMP_LIST.copy()
        TEMP_LIST.clear()
        if not call_list:
            return 0
        try:
            await Statistics.bulk_create(call_list)
        except Exception as e:
            logger.error(f"{reason}批量添加调用记录失败", "定时任务", e=e)
            retain_count = max(STATS_BUFFER_MAX_RETAIN - len(TEMP_LIST), 0)
            if retain_count:
                TEMP_LIST[:0] = call_list[-retain_count:]
            return 0
    logger.debug(f"{reason}批量添加调用记录 {len(call_list)} 条", "定时任务")
    return len(call_list)


async def _append_statistics(record: Statistics) -> None:
    TEMP_LIST.append(record)
    if len(TEMP_LIST) >= STATS_BUFFER_FLUSH_SIZE and not _STATS_FLUSH_LOCK.locked():
        await _flush_statistics_buffer("缓冲区触发")


@run_postprocessor
async def _(
    matcher: Matcher,
    exception: Exception | None,
    bot: Bot,
    session: Uninfo,
    event: Event,
):
    if matcher.type == "notice" and not isinstance(event, PokeNotifyEvent):
        """过滤除poke外的notice"""
        return
    if matcher.plugin:
        plugin = PluginInfoMemoryCache.get_by_module_path(matcher.plugin.module_name)
        if not plugin:
            # cache miss 时不查数据库，直接跳过统计，避免阻塞
            return
        if plugin.ignore_statistics:
            return
        plugin_type = plugin.plugin_type
        if plugin_type == PluginType.NORMAL:
            entity = get_entity_ids(session)
            logger.debug(f"提交调用记录: {matcher.plugin_name}...", session=session)
            await _append_statistics(
                Statistics(
                    user_id=entity.user_id,
                    group_id=entity.group_id,
                    plugin_name=matcher.plugin_name,
                    create_time=datetime.now(),
                    bot_id=bot.self_id,
                )
            )


@scheduler.scheduled_job("interval", minutes=30, max_instances=1, coalesce=True)
async def _():
    try:
        if should_pause_tasks():
            return
        await _flush_statistics_buffer("定时")
    except Exception as e:
        logger.error("定时批量添加调用记录", "定时任务", e=e)


@driver.on_shutdown
async def _flush_statistics_on_shutdown():
    await _flush_statistics_buffer("关闭")