mirror of
https://github.com/artemium428/428th-exchange-bot.git
synced 2026-09-15 16:56:20 +00:00
115 lines
4.2 KiB
Python
115 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
from dataclasses import dataclass
|
|
from threading import Lock
|
|
from typing import Optional
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CooldownDecision:
|
|
"""Result of a cooldown check before relay."""
|
|
|
|
allowed: bool
|
|
should_ack: bool
|
|
|
|
|
|
class CooldownService:
|
|
"""Per-user cooldown. In-memory; enough for a single bot process.
|
|
|
|
Cooldown is committed only after a successful relay via :meth:`commit`.
|
|
Media albums share one ack and one cooldown window.
|
|
Expired entries are pruned opportunistically so the maps stay bounded.
|
|
"""
|
|
|
|
def __init__(self, seconds: float) -> None:
|
|
self._seconds = max(0.0, seconds)
|
|
self._last_allowed: dict[int, float] = {}
|
|
# user_id -> (media_group_id, monotonic_ts when first item was accepted)
|
|
self._active_albums: dict[int, tuple[str, float]] = {}
|
|
self._services_last: dict[int, float] = {}
|
|
self._lock = Lock()
|
|
# Album items can arrive a few seconds apart.
|
|
self._album_window = max(self._seconds, 3.0)
|
|
self._services_seconds = max(self._seconds, 5.0)
|
|
|
|
def _retention(self) -> float:
|
|
return max(self._album_window, self._services_seconds, self._seconds, 1.0) * 2
|
|
|
|
def _prune_unlocked(self, now: float) -> None:
|
|
retain = self._retention()
|
|
stale_users = [uid for uid, ts in self._last_allowed.items() if now - ts > retain]
|
|
for uid in stale_users:
|
|
del self._last_allowed[uid]
|
|
|
|
stale_albums = [
|
|
uid for uid, (_gid, ts) in self._active_albums.items() if now - ts > retain
|
|
]
|
|
for uid in stale_albums:
|
|
del self._active_albums[uid]
|
|
|
|
stale_services = [
|
|
uid for uid, ts in self._services_last.items() if now - ts > retain
|
|
]
|
|
for uid in stale_services:
|
|
del self._services_last[uid]
|
|
|
|
def check(self, user_id: int, media_group_id: Optional[str] = None) -> CooldownDecision:
|
|
"""Return whether the user may relay, without consuming the cooldown."""
|
|
if self._seconds <= 0 and media_group_id is None:
|
|
return CooldownDecision(allowed=True, should_ack=True)
|
|
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
self._prune_unlocked(now)
|
|
|
|
if media_group_id is not None:
|
|
active = self._active_albums.get(user_id)
|
|
if (
|
|
active is not None
|
|
and active[0] == media_group_id
|
|
and now - active[1] < self._album_window
|
|
):
|
|
return CooldownDecision(allowed=True, should_ack=False)
|
|
|
|
last = self._last_allowed.get(user_id)
|
|
if last is not None and now - last < self._seconds:
|
|
return CooldownDecision(allowed=False, should_ack=False)
|
|
|
|
return CooldownDecision(allowed=True, should_ack=True)
|
|
|
|
def commit(self, user_id: int, media_group_id: Optional[str] = None) -> None:
|
|
"""Mark cooldown after a successful relay."""
|
|
if self._seconds <= 0 and media_group_id is None:
|
|
return
|
|
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
self._prune_unlocked(now)
|
|
|
|
if media_group_id is not None:
|
|
active = self._active_albums.get(user_id)
|
|
if (
|
|
active is not None
|
|
and active[0] == media_group_id
|
|
and now - active[1] < self._album_window
|
|
):
|
|
# Continuation of the same album — cooldown already stamped.
|
|
return
|
|
self._active_albums[user_id] = (media_group_id, now)
|
|
|
|
self._last_allowed[user_id] = now
|
|
|
|
def try_take_services(self, user_id: int) -> bool:
|
|
"""Return True once; silently False if tapped again within the services window."""
|
|
if self._services_seconds <= 0:
|
|
return True
|
|
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
self._prune_unlocked(now)
|
|
last = self._services_last.get(user_id)
|
|
if last is not None and now - last < self._services_seconds:
|
|
return False
|
|
self._services_last[user_id] = now
|
|
return True
|