mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-09-28 16:20:56 +08:00
551 lines
17 KiB
Python
551 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Iterable
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from typing import Any, Literal
|
|
|
|
from tortoise.functions import Count
|
|
|
|
from zhenxun.services.cache.bounded_ttl import BoundedTTLCache
|
|
from zhenxun.services.db_context import with_db_timeout
|
|
from zhenxun.services.message_load import is_db_unhealthy
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class GroupMemberSnapshot:
|
|
id: int
|
|
user_id: str
|
|
user_name: str
|
|
group_id: str
|
|
user_join_time: datetime | None
|
|
uid: int | None
|
|
platform: str | None
|
|
|
|
|
|
def _member_cache_sizeof(members: tuple[GroupMemberSnapshot, ...]) -> int:
|
|
size = 0
|
|
for member in members:
|
|
size += 96
|
|
size += len(member.user_id) + len(member.user_name) + len(member.group_id)
|
|
size += len(member.platform or "")
|
|
return size
|
|
|
|
|
|
_GROUP_MEMBER_CACHE = BoundedTTLCache[str, tuple[GroupMemberSnapshot, ...]](
|
|
"hot_group_info_users",
|
|
ttl_seconds=45,
|
|
max_items=512,
|
|
max_total_bytes=32 * 1024 * 1024,
|
|
sizeof=_member_cache_sizeof,
|
|
)
|
|
_GROUP_USER_IDS_CACHE = BoundedTTLCache[str, tuple[str, ...]](
|
|
"hot_group_info_user_ids",
|
|
ttl_seconds=45,
|
|
max_items=2048,
|
|
)
|
|
_GROUP_MEMBER_BY_ID_CACHE = BoundedTTLCache[str, tuple[GroupMemberSnapshot | None]](
|
|
"hot_group_info_user_by_id",
|
|
ttl_seconds=45,
|
|
max_items=50000,
|
|
)
|
|
_USER_GROUP_CACHE = BoundedTTLCache[str, tuple[str, ...]](
|
|
"hot_group_info_user_groups",
|
|
ttl_seconds=45,
|
|
max_items=4096,
|
|
)
|
|
_USER_NAME_CACHE = BoundedTTLCache[str, str](
|
|
"hot_group_info_user_names",
|
|
ttl_seconds=45,
|
|
max_items=20000,
|
|
)
|
|
_CHAT_RANK_CACHE = BoundedTTLCache[str, tuple[tuple[str, int], ...]](
|
|
"hot_chat_history_rank",
|
|
ttl_seconds=20,
|
|
max_items=512,
|
|
)
|
|
_CHAT_FIRST_MSG_CACHE = BoundedTTLCache[str, tuple[datetime | None]](
|
|
"hot_chat_history_first_msg",
|
|
ttl_seconds=300,
|
|
max_items=2048,
|
|
)
|
|
_STATISTICS_COUNT_CACHE = BoundedTTLCache[str, tuple[tuple[str, int], ...]](
|
|
"hot_statistics_plugin_counts",
|
|
ttl_seconds=20,
|
|
max_items=512,
|
|
)
|
|
|
|
_GROUP_MEMBER_LOCKS: dict[str, asyncio.Lock] = {}
|
|
_GROUP_USER_IDS_LOCKS: dict[str, asyncio.Lock] = {}
|
|
_USER_GROUP_LOCKS: dict[str, asyncio.Lock] = {}
|
|
_CHAT_RANK_LOCKS: dict[str, asyncio.Lock] = {}
|
|
_CHAT_FIRST_MSG_LOCKS: dict[str, asyncio.Lock] = {}
|
|
_STATISTICS_LOCKS: dict[str, asyncio.Lock] = {}
|
|
_MAX_LOCK_POOL_SIZE = 4096
|
|
_MEMBER_DB_TIMEOUT = 2.0
|
|
_AGGREGATE_DB_TIMEOUT = 3.0
|
|
|
|
|
|
async def _read_or_default(
|
|
coro,
|
|
*,
|
|
timeout: float,
|
|
operation: str,
|
|
default,
|
|
):
|
|
try:
|
|
return await with_db_timeout(
|
|
coro,
|
|
timeout=timeout,
|
|
operation=operation,
|
|
source="hot_query_cache",
|
|
)
|
|
except TimeoutError:
|
|
return default
|
|
|
|
|
|
def _get_lock(pool: dict[str, asyncio.Lock], key: str) -> asyncio.Lock:
|
|
lock = pool.get(key)
|
|
if lock is None:
|
|
if len(pool) >= _MAX_LOCK_POOL_SIZE:
|
|
for old_key, old_lock in list(pool.items()):
|
|
if not old_lock.locked():
|
|
pool.pop(old_key, None)
|
|
break
|
|
lock = asyncio.Lock()
|
|
pool[key] = lock
|
|
return lock
|
|
|
|
|
|
def _normalize_id(value: object) -> str:
|
|
return str(value or "")
|
|
|
|
|
|
def _normalize_ids(values: Iterable[object] | None) -> tuple[str, ...] | None:
|
|
if values is None:
|
|
return None
|
|
return tuple(dict.fromkeys(v for value in values if (v := _normalize_id(value))))
|
|
|
|
|
|
async def get_group_members(
|
|
group_id: str | int | None,
|
|
) -> tuple[GroupMemberSnapshot, ...]:
|
|
"""Return lightweight group-member snapshots with a short runtime TTL."""
|
|
group_key = _normalize_id(group_id)
|
|
if not group_key:
|
|
return ()
|
|
|
|
cached = await _GROUP_MEMBER_CACHE.get(group_key)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
lock = _get_lock(_GROUP_MEMBER_LOCKS, group_key)
|
|
async with lock:
|
|
cached = await _GROUP_MEMBER_CACHE.get(group_key)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
from zhenxun.models.group_member_info import GroupInfoUser
|
|
|
|
if is_db_unhealthy():
|
|
return ()
|
|
rows = await _read_or_default(
|
|
GroupInfoUser.filter(group_id=group_key).values_list(
|
|
"id",
|
|
"user_id",
|
|
"user_name",
|
|
"user_join_time",
|
|
"uid",
|
|
"platform",
|
|
),
|
|
timeout=_MEMBER_DB_TIMEOUT,
|
|
operation="hot_query_cache.get_group_members",
|
|
default=(),
|
|
)
|
|
members = tuple(
|
|
GroupMemberSnapshot(
|
|
id=int(row[0] or 0),
|
|
user_id=str(row[1] or ""),
|
|
user_name=str(row[2] or ""),
|
|
group_id=group_key,
|
|
user_join_time=row[3],
|
|
uid=int(row[4]) if row[4] is not None else None,
|
|
platform=str(row[5]) if row[5] else None,
|
|
)
|
|
for row in rows
|
|
if row[1]
|
|
)
|
|
await _GROUP_MEMBER_CACHE.set(group_key, members)
|
|
await _GROUP_USER_IDS_CACHE.set(
|
|
group_key, tuple(member.user_id for member in members)
|
|
)
|
|
return members
|
|
|
|
|
|
async def get_group_member_map(
|
|
group_id: str | int | None,
|
|
user_ids: Iterable[object] | None = None,
|
|
) -> dict[str, GroupMemberSnapshot]:
|
|
group_key = _normalize_id(group_id)
|
|
if not group_key:
|
|
return {}
|
|
wanted = _normalize_ids(user_ids)
|
|
if wanted is not None and not wanted:
|
|
return {}
|
|
if wanted is None:
|
|
members = await get_group_members(group_key)
|
|
return {member.user_id: member for member in members}
|
|
|
|
cached_members = await _GROUP_MEMBER_CACHE.get(group_key)
|
|
if cached_members is not None:
|
|
wanted_set = set(wanted)
|
|
return {
|
|
member.user_id: member
|
|
for member in cached_members
|
|
if member.user_id in wanted_set
|
|
}
|
|
|
|
result: dict[str, GroupMemberSnapshot] = {}
|
|
missing: list[str] = []
|
|
for user_id in wanted:
|
|
cache_key = f"{group_key}:{user_id}"
|
|
cached = await _GROUP_MEMBER_BY_ID_CACHE.get(cache_key)
|
|
if cached is None:
|
|
missing.append(user_id)
|
|
else:
|
|
member = cached[0]
|
|
if member is not None:
|
|
result[user_id] = member
|
|
|
|
if missing:
|
|
from zhenxun.models.group_member_info import GroupInfoUser
|
|
|
|
if is_db_unhealthy():
|
|
return result
|
|
rows = await _read_or_default(
|
|
GroupInfoUser.filter(group_id=group_key, user_id__in=missing).values_list(
|
|
"id",
|
|
"user_id",
|
|
"user_name",
|
|
"user_join_time",
|
|
"uid",
|
|
"platform",
|
|
),
|
|
timeout=_MEMBER_DB_TIMEOUT,
|
|
operation="hot_query_cache.get_group_member_map",
|
|
default=(),
|
|
)
|
|
found: set[str] = set()
|
|
for row in rows:
|
|
if not row[1]:
|
|
continue
|
|
member = GroupMemberSnapshot(
|
|
id=int(row[0] or 0),
|
|
user_id=str(row[1] or ""),
|
|
user_name=str(row[2] or ""),
|
|
group_id=group_key,
|
|
user_join_time=row[3],
|
|
uid=int(row[4]) if row[4] is not None else None,
|
|
platform=str(row[5]) if row[5] else None,
|
|
)
|
|
result[member.user_id] = member
|
|
found.add(member.user_id)
|
|
await _GROUP_MEMBER_BY_ID_CACHE.set(
|
|
f"{group_key}:{member.user_id}", (member,)
|
|
)
|
|
for user_id in missing:
|
|
if user_id not in found:
|
|
await _GROUP_MEMBER_BY_ID_CACHE.set(f"{group_key}:{user_id}", (None,))
|
|
return result
|
|
|
|
|
|
async def get_group_member(
|
|
group_id: str | int | None,
|
|
user_id: str | int | None,
|
|
) -> GroupMemberSnapshot | None:
|
|
user_key = _normalize_id(user_id)
|
|
if not user_key:
|
|
return None
|
|
return (await get_group_member_map(group_id, [user_key])).get(user_key)
|
|
|
|
|
|
async def get_group_user_ids(group_id: str | int | None) -> set[str]:
|
|
group_key = _normalize_id(group_id)
|
|
if not group_key:
|
|
return set()
|
|
cached = await _GROUP_USER_IDS_CACHE.get(group_key)
|
|
if cached is not None:
|
|
return set(cached)
|
|
|
|
cached_members = await _GROUP_MEMBER_CACHE.get(group_key)
|
|
if cached_members is not None:
|
|
user_ids = tuple(member.user_id for member in cached_members)
|
|
await _GROUP_USER_IDS_CACHE.set(group_key, user_ids)
|
|
return set(user_ids)
|
|
|
|
lock = _get_lock(_GROUP_USER_IDS_LOCKS, group_key)
|
|
async with lock:
|
|
cached = await _GROUP_USER_IDS_CACHE.get(group_key)
|
|
if cached is not None:
|
|
return set(cached)
|
|
|
|
from zhenxun.models.group_member_info import GroupInfoUser
|
|
|
|
if is_db_unhealthy():
|
|
return set()
|
|
rows = await _read_or_default(
|
|
GroupInfoUser.filter(group_id=group_key).values_list("user_id", flat=True),
|
|
timeout=_MEMBER_DB_TIMEOUT,
|
|
operation="hot_query_cache.get_group_user_ids",
|
|
default=(),
|
|
)
|
|
user_ids = tuple(str(user_id) for user_id in rows if user_id)
|
|
await _GROUP_USER_IDS_CACHE.set(group_key, user_ids)
|
|
return set(user_ids)
|
|
|
|
|
|
async def get_user_group_ids(user_id: str | int | None) -> list[str]:
|
|
user_key = _normalize_id(user_id)
|
|
if not user_key:
|
|
return []
|
|
|
|
cached = await _USER_GROUP_CACHE.get(user_key)
|
|
if cached is not None:
|
|
return list(cached)
|
|
|
|
lock = _get_lock(_USER_GROUP_LOCKS, user_key)
|
|
async with lock:
|
|
cached = await _USER_GROUP_CACHE.get(user_key)
|
|
if cached is not None:
|
|
return list(cached)
|
|
|
|
from zhenxun.models.group_member_info import GroupInfoUser
|
|
|
|
if is_db_unhealthy():
|
|
return []
|
|
rows = await _read_or_default(
|
|
GroupInfoUser.filter(user_id=user_key).values_list("group_id", flat=True),
|
|
timeout=_MEMBER_DB_TIMEOUT,
|
|
operation="hot_query_cache.get_user_group_ids",
|
|
default=(),
|
|
)
|
|
group_ids = tuple(str(group_id) for group_id in rows if group_id)
|
|
await _USER_GROUP_CACHE.set(user_key, group_ids)
|
|
return list(group_ids)
|
|
|
|
|
|
async def get_member_names(
|
|
user_ids: Iterable[object],
|
|
group_id: str | int | None = None,
|
|
) -> dict[str, str]:
|
|
user_keys = _normalize_ids(user_ids) or ()
|
|
if not user_keys:
|
|
return {}
|
|
if group_id:
|
|
members = await get_group_member_map(group_id, user_keys)
|
|
return {user_id: members[user_id].user_name for user_id in members}
|
|
|
|
result: dict[str, str] = {}
|
|
missing: list[str] = []
|
|
for user_id in user_keys:
|
|
cached = await _USER_NAME_CACHE.get(user_id)
|
|
if cached is None:
|
|
missing.append(user_id)
|
|
else:
|
|
result[user_id] = cached
|
|
|
|
if missing:
|
|
from zhenxun.models.group_member_info import GroupInfoUser
|
|
|
|
if is_db_unhealthy():
|
|
return result
|
|
rows = await _read_or_default(
|
|
GroupInfoUser.filter(user_id__in=missing).values_list(
|
|
"user_id", "user_name"
|
|
),
|
|
timeout=_MEMBER_DB_TIMEOUT,
|
|
operation="hot_query_cache.get_member_names",
|
|
default=(),
|
|
)
|
|
for user_id, user_name in rows:
|
|
user_key = str(user_id)
|
|
if user_key not in result:
|
|
result[user_key] = str(user_name or "")
|
|
for user_id in missing:
|
|
await _USER_NAME_CACHE.set(user_id, result.get(user_id, ""))
|
|
return result
|
|
|
|
|
|
async def get_member_name(
|
|
user_id: str | int | None,
|
|
group_id: str | int | None = None,
|
|
) -> str | None:
|
|
user_key = _normalize_id(user_id)
|
|
if not user_key:
|
|
return None
|
|
return (await get_member_names([user_key], group_id)).get(user_key) or None
|
|
|
|
|
|
async def invalidate_group_members(
|
|
group_id: str | int | None = None,
|
|
user_ids: Iterable[object] | None = None,
|
|
) -> None:
|
|
if group_id is None:
|
|
await _GROUP_MEMBER_CACHE.clear()
|
|
await _GROUP_USER_IDS_CACHE.clear()
|
|
await _GROUP_MEMBER_BY_ID_CACHE.clear()
|
|
return
|
|
group_key = _normalize_id(group_id)
|
|
await _GROUP_MEMBER_CACHE.delete(group_key)
|
|
await _GROUP_USER_IDS_CACHE.delete(group_key)
|
|
normalized_ids = _normalize_ids(user_ids)
|
|
if normalized_ids is None:
|
|
await _GROUP_MEMBER_BY_ID_CACHE.clear()
|
|
return
|
|
for user_id in normalized_ids:
|
|
await _GROUP_MEMBER_BY_ID_CACHE.delete(f"{group_key}:{user_id}")
|
|
|
|
|
|
async def invalidate_member_names(user_ids: Iterable[object] | None = None) -> None:
|
|
if user_ids is None:
|
|
await _USER_NAME_CACHE.clear()
|
|
await _USER_GROUP_CACHE.clear()
|
|
return
|
|
for user_id in _normalize_ids(user_ids) or ():
|
|
await _USER_NAME_CACHE.delete(user_id)
|
|
await _USER_GROUP_CACHE.delete(user_id)
|
|
|
|
|
|
def _datetime_key(value: datetime | None) -> str:
|
|
return value.isoformat(" ", timespec="seconds") if value else ""
|
|
|
|
|
|
def _date_scope_key(date_scope: tuple[datetime, datetime] | None) -> str:
|
|
if not date_scope:
|
|
return ""
|
|
end_bucket = int(date_scope[1].timestamp() // 20)
|
|
return f"{_datetime_key(date_scope[0])}..bucket:{end_bucket}"
|
|
|
|
|
|
async def get_chat_history_rank_cached(
|
|
model: Any,
|
|
gid: str | None,
|
|
limit: int = 10,
|
|
order: str = "DESC",
|
|
date_scope: tuple[datetime, datetime] | None = None,
|
|
) -> list[tuple[str, int]]:
|
|
key = f"{gid or '*'}:{limit}:{order}:{_date_scope_key(date_scope)}"
|
|
cached = await _CHAT_RANK_CACHE.get(key)
|
|
if cached is not None:
|
|
return list(cached)
|
|
|
|
lock = _get_lock(_CHAT_RANK_LOCKS, key)
|
|
async with lock:
|
|
cached = await _CHAT_RANK_CACHE.get(key)
|
|
if cached is not None:
|
|
return list(cached)
|
|
|
|
if is_db_unhealthy():
|
|
return []
|
|
order_prefix = "-" if order == "DESC" else ""
|
|
query: Any = model.filter(group_id=gid) if gid else model
|
|
if date_scope:
|
|
filter_scope = (
|
|
date_scope[0].isoformat(" "),
|
|
date_scope[1].isoformat(" "),
|
|
)
|
|
query = query.filter(create_time__range=filter_scope)
|
|
rows = await _read_or_default(
|
|
query.annotate(count=Count("user_id"))
|
|
.order_by(f"{order_prefix}count")
|
|
.group_by("user_id")
|
|
.limit(limit)
|
|
.values_list("user_id", "count"),
|
|
timeout=_AGGREGATE_DB_TIMEOUT,
|
|
operation="hot_query_cache.get_chat_history_rank",
|
|
default=(),
|
|
)
|
|
result = tuple((str(user_id), int(count)) for user_id, count in rows)
|
|
await _CHAT_RANK_CACHE.set(key, result)
|
|
return list(result)
|
|
|
|
|
|
async def get_chat_history_first_msg_datetime_cached(
|
|
model: Any,
|
|
group_id: str | None,
|
|
) -> datetime | None:
|
|
key = group_id or "*"
|
|
cached = await _CHAT_FIRST_MSG_CACHE.get(key)
|
|
if cached is not None:
|
|
return cached[0]
|
|
|
|
lock = _get_lock(_CHAT_FIRST_MSG_LOCKS, key)
|
|
async with lock:
|
|
cached = await _CHAT_FIRST_MSG_CACHE.get(key)
|
|
if cached is not None:
|
|
return cached[0]
|
|
|
|
if is_db_unhealthy():
|
|
return None
|
|
query: Any = model.filter(group_id=group_id) if group_id else model.all()
|
|
message = await _read_or_default(
|
|
query.order_by("create_time").first(),
|
|
timeout=_AGGREGATE_DB_TIMEOUT,
|
|
operation="hot_query_cache.get_chat_history_first_msg",
|
|
default=None,
|
|
)
|
|
result = getattr(message, "create_time", None) if message else None
|
|
await _CHAT_FIRST_MSG_CACHE.set(key, (result,))
|
|
return result
|
|
|
|
|
|
async def get_statistics_plugin_counts_cached(
|
|
scope: Literal["global", "user", "group"],
|
|
*,
|
|
plugin_name: str | None,
|
|
start_time: datetime | None,
|
|
user_id: str | None = None,
|
|
group_id: str | None = None,
|
|
) -> list[tuple[str, int]]:
|
|
key = (
|
|
f"{scope}:{plugin_name or ''}:{_datetime_key(start_time)}:"
|
|
f"{user_id or ''}:{group_id or ''}"
|
|
)
|
|
cached = await _STATISTICS_COUNT_CACHE.get(key)
|
|
if cached is not None:
|
|
return list(cached)
|
|
|
|
lock = _get_lock(_STATISTICS_LOCKS, key)
|
|
async with lock:
|
|
cached = await _STATISTICS_COUNT_CACHE.get(key)
|
|
if cached is not None:
|
|
return list(cached)
|
|
|
|
if is_db_unhealthy():
|
|
return []
|
|
from zhenxun.models.statistics import Statistics
|
|
|
|
query: Any = Statistics
|
|
if scope == "user":
|
|
query = Statistics.filter(user_id=user_id)
|
|
if group_id:
|
|
query = query.filter(group_id=group_id)
|
|
elif scope == "group":
|
|
query = Statistics.filter(group_id=group_id)
|
|
if plugin_name:
|
|
query = query.filter(plugin_name=plugin_name)
|
|
if start_time:
|
|
query = query.filter(create_time__gte=start_time)
|
|
rows = await _read_or_default(
|
|
query.annotate(count=Count("id"))
|
|
.group_by("plugin_name")
|
|
.values_list("plugin_name", "count"),
|
|
timeout=_AGGREGATE_DB_TIMEOUT,
|
|
operation="hot_query_cache.get_statistics_plugin_counts",
|
|
default=(),
|
|
)
|
|
result = tuple((str(plugin), int(count)) for plugin, count in rows)
|
|
await _STATISTICS_COUNT_CACHE.set(key, result)
|
|
return list(result)
|