82 lines
2.3 KiB
Python
82 lines
2.3 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import time
|
||
|
|
from collections import OrderedDict
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from yuxi.utils.logging_config import logger
|
||
|
|
|
||
|
|
_MAX_CACHE_SIZE = 128
|
||
|
|
_MAX_AGE_S = 3600
|
||
|
|
|
||
|
|
|
||
|
|
class SentMessageCache:
|
||
|
|
def __init__(self, max_size: int = _MAX_CACHE_SIZE, max_age_s: int = _MAX_AGE_S):
|
||
|
|
self._cache: OrderedDict[str, tuple[float, dict[str, Any]]] = OrderedDict()
|
||
|
|
self._max_size = max_size
|
||
|
|
self._max_age_s = max_age_s
|
||
|
|
|
||
|
|
def put(self, msg_id: str, metadata: dict[str, Any] | None = None) -> None:
|
||
|
|
now = time.monotonic()
|
||
|
|
if len(self._cache) >= self._max_size:
|
||
|
|
self._cache.popitem(last=False)
|
||
|
|
self._cache[msg_id] = (now, metadata or {})
|
||
|
|
|
||
|
|
def get(self, msg_id: str) -> dict[str, Any] | None:
|
||
|
|
entry = self._cache.get(msg_id)
|
||
|
|
if entry is None:
|
||
|
|
return None
|
||
|
|
ts, meta = entry
|
||
|
|
if time.monotonic() - ts > self._max_age_s:
|
||
|
|
self._cache.pop(msg_id, None)
|
||
|
|
return None
|
||
|
|
return meta
|
||
|
|
|
||
|
|
def remove(self, msg_id: str) -> None:
|
||
|
|
self._cache.pop(msg_id, None)
|
||
|
|
|
||
|
|
def update(self, msg_id: str, metadata: dict[str, Any]) -> None:
|
||
|
|
entry = self._cache.get(msg_id)
|
||
|
|
if entry is None:
|
||
|
|
return
|
||
|
|
ts, existing = entry
|
||
|
|
existing.update(metadata)
|
||
|
|
self._cache[msg_id] = (ts, existing)
|
||
|
|
|
||
|
|
def size(self) -> int:
|
||
|
|
return len(self._cache)
|
||
|
|
|
||
|
|
def clear(self) -> None:
|
||
|
|
self._cache.clear()
|
||
|
|
|
||
|
|
def expired_count(self) -> int:
|
||
|
|
now = time.monotonic()
|
||
|
|
expired = 0
|
||
|
|
for _msg_id, (ts, _) in self._cache.items():
|
||
|
|
if now - ts > self._max_age_s:
|
||
|
|
expired += 1
|
||
|
|
return expired
|
||
|
|
|
||
|
|
|
||
|
|
def track_sent_message(
|
||
|
|
cache: SentMessageCache,
|
||
|
|
msg_id: str,
|
||
|
|
chat_id: str = "",
|
||
|
|
message_type: str = "text",
|
||
|
|
) -> None:
|
||
|
|
cache.put(msg_id, {"chat_id": chat_id, "message_type": message_type, "sent_at": time.monotonic()})
|
||
|
|
logger.debug(f"Sent message cached: {msg_id}")
|
||
|
|
|
||
|
|
|
||
|
|
def update_sent_message_status(
|
||
|
|
cache: SentMessageCache,
|
||
|
|
msg_id: str,
|
||
|
|
status: str,
|
||
|
|
metadata: dict[str, Any] | None = None,
|
||
|
|
) -> None:
|
||
|
|
update_data = {"status": status}
|
||
|
|
if metadata:
|
||
|
|
update_data.update(metadata)
|
||
|
|
cache.update(msg_id, update_data)
|
||
|
|
logger.debug(f"Sent message status updated: {msg_id} -> {status}")
|