mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-09-29 00:32:06 +08:00
Copaan
5596497947
* bugfix:修复notice事件扩散问题 * 优化并发调度 * bugfix:修复签到样式 * bugfix：功能调用统计修复 * bugfix:修复私聊时功能调用统计显示已退群问题 * 提高插件适配兼容性 * 优化发送队列 * 修改权限检查设计 * 继续修改权限检查设计 * 完善权限检查设计 * 优化sqlite配置 * 优化数据库初始化 * 代码整理，无用代码清理 * bugfix:修复启动时数据库校验问题 * bugfix:修复预算裁剪过于激进问题
114 lines
3.4 KiB
Python
114 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections import deque
|
|
import contextlib
|
|
import time
|
|
from typing import TYPE_CHECKING
|
|
|
|
from zhenxun.services.log import logger
|
|
|
|
if TYPE_CHECKING:
|
|
from .bot_message_store import BotMessageStore
|
|
|
|
LOG_COMMAND = "BotMessageStore"
|
|
|
|
_BUFFER_MAX_RETAIN = 20_000
|
|
_FLUSH_TRIGGER_SIZE = 64
|
|
_FLUSH_BATCH_SIZE = 500
|
|
_FLUSH_INTERVAL_SECONDS = 5.0
|
|
_DROP_LOG_INTERVAL_SECONDS = 10.0
|
|
|
|
_buffer: deque[BotMessageStore] = deque()
|
|
_buffer_lock = asyncio.Lock()
|
|
_flush_lock = asyncio.Lock()
|
|
_flush_task: asyncio.Task[None] | None = None
|
|
_dropped = 0
|
|
_last_drop_log_at = 0.0
|
|
|
|
|
|
def _ensure_flush_task() -> None:
|
|
global _flush_task
|
|
if _flush_task is not None and not _flush_task.done():
|
|
return
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
return
|
|
_flush_task = loop.create_task(_flush_loop())
|
|
|
|
|
|
def _record_drop() -> None:
|
|
global _dropped, _last_drop_log_at
|
|
_dropped += 1
|
|
now = time.monotonic()
|
|
if now - _last_drop_log_at < _DROP_LOG_INTERVAL_SECONDS:
|
|
return
|
|
_last_drop_log_at = now
|
|
logger.warning(
|
|
f"bot_message_store buffer full, dropped {_dropped} records, "
|
|
f"backlog={len(_buffer)}",
|
|
LOG_COMMAND,
|
|
)
|
|
|
|
|
|
async def _flush_loop() -> None:
|
|
while True:
|
|
await asyncio.sleep(_FLUSH_INTERVAL_SECONDS)
|
|
try:
|
|
await flush_bot_message_store_buffer("定时")
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
logger.warning("定时批量写入 Bot 发送记录失败", LOG_COMMAND, e=exc)
|
|
|
|
|
|
async def append_bot_message_store_record(record: BotMessageStore) -> None:
|
|
_ensure_flush_task()
|
|
async with _buffer_lock:
|
|
if len(_buffer) >= _BUFFER_MAX_RETAIN:
|
|
_buffer.popleft()
|
|
_record_drop()
|
|
_buffer.append(record)
|
|
should_flush = len(_buffer) >= _FLUSH_TRIGGER_SIZE and not _flush_lock.locked()
|
|
if should_flush:
|
|
await flush_bot_message_store_buffer("缓冲区触发")
|
|
|
|
|
|
async def flush_bot_message_store_buffer(reason: str) -> int:
|
|
from .bot_message_store import BotMessageStore
|
|
|
|
async with _flush_lock:
|
|
written = 0
|
|
while True:
|
|
batch: list[BotMessageStore] = []
|
|
async with _buffer_lock:
|
|
while _buffer and len(batch) < _FLUSH_BATCH_SIZE:
|
|
batch.append(_buffer.popleft())
|
|
if not batch:
|
|
break
|
|
try:
|
|
await BotMessageStore.bulk_create(batch, batch_size=_FLUSH_BATCH_SIZE)
|
|
except Exception as exc:
|
|
async with _buffer_lock:
|
|
retain_count = max(_BUFFER_MAX_RETAIN - len(_buffer), 0)
|
|
for record in reversed(batch[-retain_count:]):
|
|
_buffer.appendleft(record)
|
|
logger.error(f"{reason}批量写入 Bot 发送记录失败", LOG_COMMAND, e=exc)
|
|
return written
|
|
written += len(batch)
|
|
if written:
|
|
logger.debug(f"{reason}批量写入 Bot 发送记录 {written} 条", LOG_COMMAND)
|
|
return written
|
|
|
|
|
|
async def stop_bot_message_store_buffer() -> int:
|
|
global _flush_task
|
|
task = _flush_task
|
|
_flush_task = None
|
|
if task is not None:
|
|
task.cancel()
|
|
with contextlib.suppress(BaseException):
|
|
await task
|
|
return await flush_bot_message_store_buffer("关闭")
|