"""Hot-path account selector — pluggable strategies. Two fully independent strategies: * ``_quota_select`` — scores candidates by health / quota / inflight / fails. Used when ``account.refresh.enabled=true``. Behaviour is the historical one, kept unchanged. * ``_random_select`` — uniform random choice among non-cooling candidates. Used when ``account.refresh.enabled=false``. Ignores quota and health entirely. Strategy selection is process-global, registered once by the lifespan via :func:`set_strategy`. Callers (``AccountDirectory``) invoke :func:`select` / :func:`select_any` and are unaware of which strategy is active. """ import array # noqa: F401 — used in forward-referenced type annotations import random from typing import Literal from app.platform.config.snapshot import get_config from ..shared.enums import PoolId from .table import AccountRuntimeTable # Scoring weights used by the quota strategy. _W_HEALTH = 100.0 _W_QUOTA = 25.0 _W_RECENT = 15.0 # penalty for recently used accounts _W_INFLIGHT = 20.0 _W_FAIL = 4.0 _RECENT_WINDOW_S = 15 # seconds # --------------------------------------------------------------------------- # Strategy registry # --------------------------------------------------------------------------- _StrategyName = Literal["quota", "random"] _STRATEGY_NAME: _StrategyName = "random" def set_strategy(name: _StrategyName) -> None: """Register the process-global selection strategy. Called once by the lifespan after reading ``account.refresh.enabled``. """ global _STRATEGY_NAME if name not in ("quota", "random"): raise ValueError(f"unknown selection strategy: {name!r}") _STRATEGY_NAME = name def current_strategy() -> _StrategyName: return _STRATEGY_NAME # --------------------------------------------------------------------------- # Public entry points — delegate to the active strategy # --------------------------------------------------------------------------- def select( table: AccountRuntimeTable, pool_id: int, mode_id: int, *, exclude_idxs: frozenset[int] | None = None, prefer_tag_idxs: set[int] | None = None, now_s: int, ) -> int | None: """Select an account slot for ``(pool_id, mode_id)``. Returns the slot index or ``None`` when no candidate is available. Does not mutate the table — callers increment inflight separately. """ if _STRATEGY_NAME == "random": return _random_select( table, pool_id, exclude_idxs=exclude_idxs, prefer_tag_idxs=prefer_tag_idxs, now_s=now_s, ) return _quota_select( table, pool_id, mode_id, exclude_idxs=exclude_idxs, prefer_tag_idxs=prefer_tag_idxs, now_s=now_s, ) def select_any( table: AccountRuntimeTable, pool_id: int, *, exclude_idxs: frozenset[int] | None = None, prefer_tag_idxs: set[int] | None = None, now_s: int, ) -> int | None: """Select any active account in ``pool_id`` irrespective of per-mode quota. Used by WebSocket-based products that manage their own rate limiting. """ if _STRATEGY_NAME == "random": return _random_select( table, pool_id, exclude_idxs=exclude_idxs, prefer_tag_idxs=prefer_tag_idxs, now_s=now_s, ) return _quota_select_any( table, pool_id, exclude_idxs=exclude_idxs, prefer_tag_idxs=prefer_tag_idxs, now_s=now_s, ) # --------------------------------------------------------------------------- # Strategy: quota — score-based selection (unchanged behaviour) # --------------------------------------------------------------------------- def _quota_select( table: AccountRuntimeTable, pool_id: int, mode_id: int, *, exclude_idxs: frozenset[int] | None, prefer_tag_idxs: set[int] | None, now_s: int, ) -> int | None: candidates: set[int] | None = table.mode_available.get((pool_id, mode_id)) if not candidates: return None reset_col = table._reset_col(mode_id) quota_col = table._quota_col(mode_id) total_col = table._total_col(mode_id) window_col = table._window_col(mode_id) _maybe_reset_windows( table, candidates, mode_id, reset_col, quota_col, total_col, window_col, pool_id, now_s, ) working: set[int] = candidates.copy() if exclude_idxs: working -= exclude_idxs working = {idx for idx in working if int(quota_col[idx]) > 0} if not working: return None if prefer_tag_idxs: preferred = working & prefer_tag_idxs working = preferred if preferred else working return _best(table, working, quota_col, now_s) def _quota_select_any( table: AccountRuntimeTable, pool_id: int, *, exclude_idxs: frozenset[int] | None, prefer_tag_idxs: set[int] | None, now_s: int, ) -> int | None: candidates: set[int] = _pool_union(table, pool_id) if not candidates: return None working = candidates.copy() if exclude_idxs: working -= exclude_idxs if not working: return None if prefer_tag_idxs: preferred = working & prefer_tag_idxs working = preferred if preferred else working return _best_no_quota(table, working, now_s) def _maybe_reset_windows( table: AccountRuntimeTable, candidates: set[int], mode_id: int, reset_col: "array.array", quota_col: "array.array", total_col: "array.array", window_col: "array.array", pool_id: int, now_s: int, ) -> None: """Reset expired windows for basic-pool accounts inline (no API call needed).""" if pool_id != int(PoolId.BASIC): return for idx in list(candidates): r = reset_col[idx] if r == 0 or now_s < r: continue if int(table.pool_by_idx[idx]) != pool_id: continue new_total = int(total_col[idx]) window_s = int(window_col[idx]) if new_total <= 0 or window_s <= 0: continue quota_col[idx] = new_total reset_col[idx] = now_s + window_s def _best( table: AccountRuntimeTable, working: set[int], quota_col: "array.array", now_s: int, ) -> int | None: best_idx = -1 best_score = -1e18 health_col = table.health_by_idx inflight_col = table.inflight_by_idx fail_col = table.fail_count_by_idx last_use_col = table.last_use_at_by_idx for idx in working: quota = int(quota_col[idx]) if quota <= 0: continue health = float(health_col[idx]) inflight = int(inflight_col[idx]) fails = min(int(fail_col[idx]), 10) last_use = int(last_use_col[idx]) score = ( health * _W_HEALTH + quota * _W_QUOTA - inflight * _W_INFLIGHT - fails * _W_FAIL ) if last_use > 0: age_s = now_s - last_use if age_s < _RECENT_WINDOW_S: score -= (1.0 - age_s / _RECENT_WINDOW_S) * _W_RECENT if score > best_score: best_score = score best_idx = idx return best_idx if best_idx >= 0 else None def _best_no_quota( table: AccountRuntimeTable, working: set[int], now_s: int, ) -> int | None: best_idx = -1 best_score = -1e18 health_col = table.health_by_idx inflight_col = table.inflight_by_idx fail_col = table.fail_count_by_idx last_use_col = table.last_use_at_by_idx for idx in working: health = float(health_col[idx]) inflight = int(inflight_col[idx]) fails = min(int(fail_col[idx]), 10) last_use = int(last_use_col[idx]) score = health * _W_HEALTH - inflight * _W_INFLIGHT - fails * _W_FAIL if last_use > 0: age_s = now_s - last_use if age_s < _RECENT_WINDOW_S: score -= (1.0 - age_s / _RECENT_WINDOW_S) * _W_RECENT if score > best_score: best_score = score best_idx = idx return best_idx if best_idx >= 0 else None # --------------------------------------------------------------------------- # Strategy: random — uniform choice with cooling + inflight filter # --------------------------------------------------------------------------- def _random_select( table: AccountRuntimeTable, pool_id: int, *, exclude_idxs: frozenset[int] | None, prefer_tag_idxs: set[int] | None, now_s: int, ) -> int | None: candidates: set[int] = _pool_union(table, pool_id) if not candidates: return None max_inflight = int(get_config("account.selection.max_inflight", 8)) cooling_col = table.cooling_until_s_by_idx inflight_col = table.inflight_by_idx working = candidates.copy() if exclude_idxs: working -= exclude_idxs working = { idx for idx in working if int(cooling_col[idx]) <= now_s and int(inflight_col[idx]) < max_inflight } if not working: return None if prefer_tag_idxs: preferred = working & prefer_tag_idxs working = preferred if preferred else working return random.choice(tuple(working)) # --------------------------------------------------------------------------- # Shared helper # --------------------------------------------------------------------------- def _pool_union(table: AccountRuntimeTable, pool_id: int) -> set[int]: """Union of all ``mode_available`` buckets for ``pool_id``.""" out: set[int] = set() for (pid, _mid), accounts in table.mode_available.items(): if pid == pool_id: out |= accounts return out __all__ = ["select", "select_any", "set_strategy", "current_strategy"]