fix: delay quota window start until first use

This commit is contained in:
JOJO 2025-11-24 01:38:24 +08:00
parent fd797e3e36
commit 7c7038b5c9

View File

@ -90,35 +90,58 @@ class UsageTracker:
return QUOTA_DEFAULTS["search_daily"] return QUOTA_DEFAULTS["search_daily"]
return QUOTA_DEFAULTS["default"][metric] return QUOTA_DEFAULTS["default"][metric]
def _ensure_window(self, metric: QuotaKey) -> Tuple[int, Optional[str], datetime, datetime]: def _ensure_window(
self,
metric: QuotaKey,
init_if_missing: bool = False,
) -> Tuple[int, Optional[str], Optional[datetime], Optional[datetime]]:
"""Returns tuple(count, window_start_iso, window_start_dt, reset_at_dt).""" """Returns tuple(count, window_start_iso, window_start_dt, reset_at_dt)."""
config = self._get_quota_config(metric) config = self._get_quota_config(metric)
window_hours = config["window_hours"] window_hours = config["window_hours"]
window_delta = timedelta(hours=window_hours) window_delta = timedelta(hours=window_hours)
window_data = self._state["windows"].setdefault(metric, {"count": 0, "window_start": None}) window_data = self._state["windows"].setdefault(metric, {"count": 0, "window_start": None})
count = int(window_data.get("count", 0))
window_start_iso = window_data.get("window_start") window_start_iso = window_data.get("window_start")
now = datetime.utcnow() now = datetime.utcnow()
window_start_dt: Optional[datetime] = None
if window_start_iso: if window_start_iso:
try: try:
parsed = datetime.fromisoformat(window_start_iso.replace("Z", "")) parsed = datetime.fromisoformat(window_start_iso.replace("Z", ""))
window_start_dt = floor_to_hour(parsed) window_start_dt = floor_to_hour(parsed)
except ValueError: except ValueError:
window_start_dt = floor_to_hour(now) window_start_dt = None
else:
window_start_dt = floor_to_hour(now)
window_data["window_start"] = frame_iso(window_start_dt)
reset_at_dt: Optional[datetime] = None
if window_start_dt:
reset_at_dt = window_start_dt + window_delta reset_at_dt = window_start_dt + window_delta
if now >= reset_at_dt and window_data.get("window_start"): if now >= reset_at_dt:
window_data["count"] = 0 window_data["count"] = 0
count = 0
if init_if_missing:
window_start_dt = floor_to_hour(now) window_start_dt = floor_to_hour(now)
reset_at_dt = window_start_dt + window_delta
window_data["window_start"] = frame_iso(window_start_dt) window_data["window_start"] = frame_iso(window_start_dt)
reset_at_dt = window_start_dt + window_delta
else:
window_start_dt = None
window_data["window_start"] = None
reset_at_dt = None
if window_start_dt is None and init_if_missing:
window_start_dt = floor_to_hour(now)
window_data["window_start"] = frame_iso(window_start_dt)
reset_at_dt = window_start_dt + window_delta
elif count == 0 and window_data.get("window_start") and not init_if_missing:
# 没有真实用量时清理遗留的窗口起点,使下一次调用时重新计窗。
window_data["window_start"] = None
window_start_iso = None
if window_start_dt and not reset_at_dt:
reset_at_dt = window_start_dt + window_delta
return ( return (
int(window_data.get("count", 0)), count,
window_data["window_start"], window_data.get("window_start"),
window_start_dt, window_start_dt,
reset_at_dt, reset_at_dt,
) )
@ -126,7 +149,9 @@ class UsageTracker:
def check_and_increment(self, metric: QuotaKey) -> Tuple[bool, Dict[str, str]]: def check_and_increment(self, metric: QuotaKey) -> Tuple[bool, Dict[str, str]]:
"""Check quota and increment if allowed. Returns (allowed, info).""" """Check quota and increment if allowed. Returns (allowed, info)."""
with self._lock: with self._lock:
count, window_start_iso, window_start_dt, reset_at_dt = self._ensure_window(metric) count, window_start_iso, window_start_dt, reset_at_dt = self._ensure_window(
metric, init_if_missing=True
)
quota = self._get_quota_config(metric)["limit"] quota = self._get_quota_config(metric)["limit"]
if count >= quota: if count >= quota:
reset_at_iso = frame_iso(reset_at_dt) reset_at_iso = frame_iso(reset_at_dt)
@ -139,8 +164,6 @@ class UsageTracker:
new_count = count + 1 new_count = count + 1
self._state["windows"][metric]["count"] = new_count self._state["windows"][metric]["count"] = new_count
if not window_start_iso:
self._state["windows"][metric]["window_start"] = frame_iso(window_start_dt)
self._state["updated_at"] = self._now_iso() self._state["updated_at"] = self._now_iso()
self._save() self._save()