"""Process-wide traffic control for hosted Douyin accounts. The backend normally runs one asyncio event loop. A controller is kept per event loop so unit tests that create a fresh loop for every test do not reuse asyncio primitives bound to an older loop. """ from __future__ import annotations import asyncio import contextvars import logging import os import time import weakref from collections import deque from contextlib import asynccontextmanager from dataclasses import dataclass from typing import Any, Awaitable, Callable, Optional, TypeVar logger = logging.getLogger("douyin_im.traffic") T = TypeVar("T") def _env_float(name: str, default: float, minimum: float = 0.0) -> float: try: return max(minimum, float(os.getenv(name, str(default)))) except (TypeError, ValueError): return default def _env_int(name: str, default: int, minimum: int = 1) -> int: try: return max(minimum, int(os.getenv(name, str(default)))) except (TypeError, ValueError): return default @dataclass(eq=False) class _SendJob: account_id: int operation: Callable[[], Awaitable[Any]] future: asyncio.Future description: str queued_at: float cancelled: bool = False operation_task: Optional[asyncio.Task] = None class GlobalSendQueue: """One paced, account-fair dispatcher for every outbound IM write. Each account owns a FIFO. After one job is selected, that account goes to the back of the rotation. A busy account therefore cannot starve smaller accounts, while messages within one account retain their original order. """ def __init__(self, interval_seconds: float = 1.0) -> None: self.interval_seconds = max(0.0, float(interval_seconds or 0.0)) self._queues: dict[int, deque[_SendJob]] = {} self._rotation: deque[int] = deque() self._in_rotation: set[int] = set() self._active_job: Optional[_SendJob] = None self._active_account: Optional[int] = None self._task: Optional[asyncio.Task] = None self._running = False self._state_lock = asyncio.Lock() self._wake = asyncio.Event() self._next_allowed_at = 0.0 @property def queued_count(self) -> int: return sum(len(queue) for queue in self._queues.values()) @property def pending_count(self) -> int: return self.queued_count + (1 if self._active_job is not None else 0) async def start(self) -> None: async with self._state_lock: if self._task and not self._task.done(): return self._running = True self._wake.clear() self._task = asyncio.create_task( self._run(), name="douyin-global-send-queue", ) async def submit( self, account_id: int, operation: Callable[[], Awaitable[T]], description: str = "", ) -> T: await self.start() loop = asyncio.get_running_loop() account_key = int(account_id or 0) future: asyncio.Future = loop.create_future() job = _SendJob( account_id=account_key, operation=operation, future=future, description=str(description or "send"), queued_at=time.time(), ) async with self._state_lock: queue = self._queues.setdefault(account_key, deque()) queue.append(job) if account_key != self._active_account and account_key not in self._in_rotation: self._rotation.append(account_key) self._in_rotation.add(account_key) position = self.pending_count self._wake.set() if position > 1: logger.info( "Outbound queued account=%s position=%s pending=%s (%s)", account_key, position, position, job.description, ) try: # Shield lets us remove/cancel the exact queued job when its caller # is cancelled instead of asyncio cancelling only the result future. return await asyncio.shield(future) except asyncio.CancelledError: active = await self._cancel_job(job) if active: # Once the network operation has begun, return its real result # instead of reporting a timeout and then delivering later. # This also keeps the single lane occupied until a to_thread # upload has truly finished. return await asyncio.shield(future) raise async def _cancel_job(self, job: _SendJob) -> bool: async with self._state_lock: if job.cancelled: return bool(job.operation_task and not job.operation_task.done()) if ( self._active_job is job and job.operation_task is not None ): # An active image upload may be running in asyncio.to_thread(). # Cancelling its awaiter does not stop the thread and would let # the next send overlap it. Let the active job drain and make # the cancelled caller wait for its authoritative result. return True job.cancelled = True queue = self._queues.get(job.account_id) if queue: try: queue.remove(job) except ValueError: pass if not queue: self._queues.pop(job.account_id, None) self._remove_from_rotation(job.account_id) if not job.future.done(): job.future.cancel() self._wake.set() return False def _remove_from_rotation(self, account_id: int) -> None: if account_id not in self._in_rotation: return self._in_rotation.discard(account_id) try: self._rotation.remove(account_id) except ValueError: pass def _pop_next_locked(self) -> Optional[_SendJob]: while self._rotation: account_id = self._rotation.popleft() self._in_rotation.discard(account_id) queue = self._queues.get(account_id) if not queue: self._queues.pop(account_id, None) continue while queue and (queue[0].cancelled or queue[0].future.cancelled()): queue.popleft() if not queue: self._queues.pop(account_id, None) continue job = queue.popleft() if not queue: self._queues.pop(account_id, None) self._active_job = job self._active_account = account_id return job return None async def _finish_job(self, job: _SendJob) -> None: async with self._state_lock: if self._active_job is job: self._active_job = None self._active_account = None queue = self._queues.get(job.account_id) while queue and (queue[0].cancelled or queue[0].future.cancelled()): queue.popleft() if queue: if job.account_id not in self._in_rotation: # Reinsert only after the active operation finishes. New # accounts that arrived while it ran get their turn first. self._rotation.append(job.account_id) self._in_rotation.add(job.account_id) else: self._queues.pop(job.account_id, None) self._remove_from_rotation(job.account_id) self._wake.set() async def _run(self) -> None: try: while True: async with self._state_lock: if not self._running: return job = self._pop_next_locked() if job is None: self._wake.clear() if job is None: await self._wake.wait() continue try: remaining = self._next_allowed_at - asyncio.get_running_loop().time() if remaining > 0: await asyncio.sleep(remaining) if job.cancelled or job.future.cancelled(): continue operation_task = asyncio.create_task(job.operation()) async with self._state_lock: job.operation_task = operation_task cancelled_before_start = job.cancelled or job.future.cancelled() if cancelled_before_start: operation_task.cancel() try: result = await operation_task except asyncio.CancelledError: if not job.future.done(): job.future.cancel() if not self._running: raise except Exception as exc: if not job.future.done(): job.future.set_exception(exc) else: if not job.future.done(): job.future.set_result(result) finally: self._next_allowed_at = ( asyncio.get_running_loop().time() + self.interval_seconds ) await self._finish_job(job) except asyncio.CancelledError: raise async def snapshot(self) -> dict[str, Any]: async with self._state_lock: per_account = { account_id: len(queue) for account_id, queue in self._queues.items() if queue } return { "running": bool(self._task and not self._task.done()), "pending_count": sum(per_account.values()) + (1 if self._active_job is not None else 0), "queued_count": sum(per_account.values()), "active_account_id": self._active_account, "active_description": ( self._active_job.description if self._active_job else "" ), "interval_seconds": self.interval_seconds, "per_account": per_account, } async def cancel_account(self, account_id: int) -> int: """Cancel callers for one stopped account without overlapping traffic. Waiting jobs are removed. An already-running operation is allowed to drain inside the dispatcher (notably, Python cannot stop a running to_thread image upload), but its caller is released immediately and no later job can start until that drain finishes. """ account_key = int(account_id or 0) cancelled = 0 async with self._state_lock: queue = self._queues.pop(account_key, deque()) self._remove_from_rotation(account_key) for job in queue: job.cancelled = True if not job.future.done(): job.future.cancel() cancelled += 1 if self._active_job and self._active_job.account_id == account_key: self._active_job.cancelled = True if not self._active_job.future.done(): self._active_job.future.cancel() cancelled += 1 self._wake.set() return cancelled async def stop(self) -> None: drain_timeout = _env_float("KEFU_SEND_SHUTDOWN_DRAIN_SECONDS", 30.0) async with self._state_lock: self._running = False task = self._task self._task = None active_job = self._active_job active_task = active_job.operation_task if active_job else None futures = [ job.future for queue in self._queues.values() for job in queue ] for queue in self._queues.values(): for job in queue: job.cancelled = True if active_job: active_job.cancelled = True futures.append(active_job.future) self._queues.clear() self._rotation.clear() self._in_rotation.clear() self._wake.set() for future in futures: if not future.done(): future.cancel() if active_task and not active_task.done() and drain_timeout > 0: try: await asyncio.wait_for( asyncio.shield(active_task), timeout=drain_timeout, ) except asyncio.TimeoutError: logger.warning( "Active outbound task did not drain within %.1fs during shutdown (%s)", drain_timeout, active_job.description if active_job else "send", ) except asyncio.CancelledError: pass except Exception: # The dispatcher owns normal operation error reporting/results. pass if active_task and not active_task.done(): active_task.cancel() if task and task is not asyncio.current_task(): if not task.done(): try: await asyncio.wait_for(asyncio.shield(task), timeout=1.0) except (asyncio.TimeoutError, asyncio.CancelledError): task.cancel() if not task.done(): try: await task except asyncio.CancelledError: pass async with self._state_lock: self._active_job = None self._active_account = None self._next_allowed_at = 0.0 self._wake.clear() class TrafficController: """Independent lanes for message writes and lower-priority network work.""" def __init__(self) -> None: self.send_queue = GlobalSendQueue( interval_seconds=_env_float( "KEFU_GLOBAL_SEND_INTERVAL_SECONDS", 1.0, ) ) background_capacity = _env_int("KEFU_BACKGROUND_NETWORK_CONCURRENCY", 4) self._background = asyncio.Semaphore(background_capacity) # Recurring polls can create hundreds of waiters when many accounts # are online. Admit at most one normal waiter to the semaphore at a # time so startup validation can join near the front instead of being # buried behind the entire polling backlog. The shared semaphore is # still the single bandwidth cap; startup work does not add extra # network concurrency. self._background_normal_admission = asyncio.Lock() self._background_startup_clear = asyncio.Event() self._background_startup_clear.set() # Starting several hundred accounts keeps startup waiters queued for # many minutes on end. Leave one slot of the shared lane for recurring # work so hosted accounts keep receiving messages during a bulk start # instead of going silent until the last account is up. self._background_startup = asyncio.Semaphore( max(1, background_capacity - 1) ) self._background_normal_max_defer = _env_float( "KEFU_BACKGROUND_NORMAL_MAX_DEFER_SECONDS", 5.0, ) self._browser = asyncio.Semaphore( _env_int("KEFU_BROWSER_START_CONCURRENCY", 1) ) self._media_proxy = asyncio.Semaphore( _env_int("KEFU_MEDIA_PROXY_CONCURRENCY", 4) ) self._background_owner: contextvars.ContextVar[tuple[Optional[asyncio.Task], int]] = ( contextvars.ContextVar( f"douyin_background_owner_{id(self)}", default=(None, 0), ) ) self.background_waiting = 0 self.background_active = 0 self.background_startup_waiting = 0 self.background_startup_active = 0 self.browser_waiting = 0 self.browser_active = 0 self.media_proxy_active = 0 @asynccontextmanager async def background_slot( self, account_id: int = 0, description: str = "request", *, startup: bool = False, ): current_task = asyncio.current_task() owner_task, depth = self._background_owner.get() if owner_task is current_task and depth > 0: token = self._background_owner.set((current_task, depth + 1)) try: yield finally: self._background_owner.reset(token) return started = asyncio.get_running_loop().time() self.background_waiting += 1 try: if startup: self.background_startup_waiting += 1 self._background_startup_clear.clear() startup_reservation = False try: await self._background_startup.acquire() startup_reservation = True await self._background.acquire() except BaseException: if startup_reservation: self._background_startup.release() raise finally: self.background_startup_waiting -= 1 if self.background_startup_waiting == 0: self._background_startup_clear.set() else: # Only one recurring/background request may wait directly on # the shared semaphore. A later startup request therefore # has at most one normal request ahead of it, not hundreds. async with self._background_normal_admission: # Yield to pending startup work, but only while a startup # waiter could still claim a slot, and never for longer # than the deferral budget. A batch of several hundred # accounts otherwise keeps startup waiters pending for the # whole run, which stalled every recurring poll behind it. if ( self._background_normal_max_defer > 0 and not self._background_startup_clear.is_set() and not self._background_startup.locked() ): try: await asyncio.wait_for( self._background_startup_clear.wait(), timeout=self._background_normal_max_defer, ) except asyncio.TimeoutError: pass await self._background.acquire() except BaseException: self.background_waiting -= 1 raise self.background_waiting -= 1 self.background_active += 1 if startup: self.background_startup_active += 1 token = self._background_owner.set((current_task, 1)) waited = asyncio.get_running_loop().time() - started if waited >= 1.0: logger.info( "Background request dequeued account=%s waited=%.2fs (%s)", account_id, waited, description, ) try: yield finally: self._background_owner.reset(token) if startup: self.background_startup_active -= 1 self._background_startup.release() self.background_active -= 1 self._background.release() @asynccontextmanager async def browser_slot(self, account_id: int = 0, description: str = "browser-login"): started = asyncio.get_running_loop().time() self.browser_waiting += 1 try: await self._browser.acquire() except BaseException: self.browser_waiting -= 1 raise self.browser_waiting -= 1 self.browser_active += 1 waited = asyncio.get_running_loop().time() - started if waited >= 1.0: logger.info( "Browser task dequeued account=%s waited=%.2fs (%s)", account_id, waited, description, ) try: yield finally: self.browser_active -= 1 self._browser.release() @asynccontextmanager async def media_proxy_slot(self): await self._media_proxy.acquire() self.media_proxy_active += 1 try: yield finally: self.media_proxy_active -= 1 self._media_proxy.release() async def snapshot(self) -> dict[str, Any]: return { "send": await self.send_queue.snapshot(), "background": { "active": self.background_active, "waiting": self.background_waiting, "startup_active": self.background_startup_active, "startup_waiting": self.background_startup_waiting, }, "browser": { "active": self.browser_active, "waiting": self.browser_waiting, }, "media_proxy": {"active": self.media_proxy_active}, } async def stop(self) -> None: await self.send_queue.stop() _CONTROLLERS: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, TrafficController]" = ( weakref.WeakKeyDictionary() ) def get_traffic_controller() -> TrafficController: loop = asyncio.get_running_loop() controller = _CONTROLLERS.get(loop) if controller is None: controller = TrafficController() _CONTROLLERS[loop] = controller return controller async def submit_outbound( account_id: int, operation: Callable[[], Awaitable[T]], description: str = "", ) -> T: return await get_traffic_controller().send_queue.submit( account_id, operation, description=description, ) async def shutdown_traffic_controller() -> None: loop = asyncio.get_running_loop() controller = _CONTROLLERS.pop(loop, None) if controller is not None: await controller.stop()