更新
This commit is contained in:
Your Name
@@ -0,0 +1,537 @@
|
||||
"""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,
|
||||
)
|
||||
)
|
||||
self._background = asyncio.Semaphore(
|
||||
_env_int("KEFU_BACKGROUND_NETWORK_CONCURRENCY", 2)
|
||||
)
|
||||
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.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"):
|
||||
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:
|
||||
await self._background.acquire()
|
||||
except BaseException:
|
||||
self.background_waiting -= 1
|
||||
raise
|
||||
self.background_waiting -= 1
|
||||
self.background_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)
|
||||
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,
|
||||
},
|
||||
"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()
|
||||
Reference in New Issue
Block a user