610 lines
22 KiB
Python
610 lines
22 KiB
Python
"""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()
|