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_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.db_context import with_db_timeout
from zhenxun.services.log import logger
from zhenxun.services.low_priority_writer import (
    LowPriorityWriterConfig,
    append_low_priority_record,
    flush_low_priority_writer,
    register_low_priority_writer,
)
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
_STATS_FLUSH_LOCK = asyncio.Lock()
_WRITER_NAME = "statistics"
driver = get_driver()


async def _write_statistics_batch(batch: list[Statistics], reason: str) -> None:
    await with_db_timeout(
        Statistics.bulk_create(batch),
        timeout=5.0,
        operation=f"Statistics.bulk_create[{len(batch)}]",
        source=f"statistics:{reason}",
    )


register_low_priority_writer(
    LowPriorityWriterConfig(
        name=_WRITER_NAME,
        write_batch=_write_statistics_batch,
        batch_size=STATS_BUFFER_FLUSH_SIZE,
        trigger_size=STATS_BUFFER_FLUSH_SIZE,
        max_retain=STATS_BUFFER_MAX_RETAIN,
        flush_interval_seconds=30 * 60,
        max_items_per_cycle=STATS_BUFFER_FLUSH_SIZE,
        backoff_base_seconds=30.0,
        backoff_max_seconds=600.0,
        log_command="定时任务",
    )
)


async def _flush_statistics_buffer(reason: str) -> int:
    """Compatibility entry used by memory governor and shutdown hooks."""
    async with _STATS_FLUSH_LOCK:
        return await flush_low_priority_writer(_WRITER_NAME, reason, force=True)


async def _append_statistics(record: Statistics) -> None:
    await append_low_priority_record(_WRITER_NAME, record)


@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,
                )
            )


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