Files
dy/backend/tests/test_traffic_control.py
T
2026-07-30 10:06:53 +08:00

641 lines
23 KiB
Python

import asyncio
import os
import sys
import unittest
from pathlib import Path
from unittest.mock import patch
BACKEND_DIR = Path(__file__).resolve().parents[1]
if str(BACKEND_DIR) not in sys.path:
sys.path.insert(0, str(BACKEND_DIR))
from rpa_engine.douyin_im.traffic_control import (
GlobalSendQueue,
TrafficController,
get_traffic_controller,
)
class GlobalSendQueueTests(unittest.IsolatedAsyncioTestCase):
async def _make_queue(self, interval_seconds: float = 0.0) -> GlobalSendQueue:
queue = GlobalSendQueue(interval_seconds=interval_seconds)
self.addAsyncCleanup(queue.stop)
return queue
async def _wait_for_queued_count(
self,
queue: GlobalSendQueue,
expected: int,
timeout: float = 0.5,
) -> None:
loop = asyncio.get_running_loop()
deadline = loop.time() + timeout
while loop.time() < deadline:
if (await queue.snapshot())["queued_count"] == expected:
return
await asyncio.sleep(0.001)
self.assertEqual((await queue.snapshot())["queued_count"], expected)
async def test_global_send_concurrency_is_one_across_accounts(self):
queue = await self._make_queue()
active = 0
maximum_active = 0
async def operation(index: int) -> int:
nonlocal active, maximum_active
active += 1
maximum_active = max(maximum_active, active)
await asyncio.sleep(0.015)
active -= 1
return index
tasks = [
asyncio.create_task(
queue.submit(
account_id=(index % 3) + 1,
operation=lambda index=index: operation(index),
description=f"send-{index}",
)
)
for index in range(9)
]
results = await asyncio.wait_for(asyncio.gather(*tasks), timeout=1.0)
self.assertEqual(sorted(results), list(range(9)))
self.assertEqual(maximum_active, 1)
self.assertEqual(active, 0)
async def test_same_account_keeps_fifo_order(self):
queue = await self._make_queue()
calls: list[int] = []
async def operation(index: int) -> int:
calls.append(index)
await asyncio.sleep(0)
return index
tasks = [
asyncio.create_task(
queue.submit(
account_id=101,
operation=lambda index=index: operation(index),
description=str(index),
)
)
for index in range(6)
]
results = await asyncio.wait_for(asyncio.gather(*tasks), timeout=0.5)
self.assertEqual(calls, list(range(6)))
self.assertEqual(results, list(range(6)))
async def test_busy_accounts_are_dispatched_round_robin(self):
queue = await self._make_queue()
calls: list[str] = []
first_started = asyncio.Event()
release_first = asyncio.Event()
async def first_operation() -> str:
calls.append("a1")
first_started.set()
await release_first.wait()
return "a1"
async def operation(name: str) -> str:
calls.append(name)
await asyncio.sleep(0)
return name
first = asyncio.create_task(queue.submit(1, first_operation, "a1"))
await asyncio.wait_for(first_started.wait(), timeout=0.2)
# Account 1 is deliberately noisy. Accounts 2 and 3 must each get a
# turn before account 1 is allowed to dispatch its next waiting job.
remaining_specs = [
(1, "a2"),
(1, "a3"),
(2, "b1"),
(2, "b2"),
(3, "c1"),
]
remaining = [
asyncio.create_task(
queue.submit(
account_id,
lambda name=name: operation(name),
description=name,
)
)
for account_id, name in remaining_specs
]
await self._wait_for_queued_count(queue, expected=len(remaining))
release_first.set()
results = await asyncio.wait_for(
asyncio.gather(first, *remaining),
timeout=0.5,
)
self.assertEqual(calls, ["a1", "b1", "c1", "a2", "b2", "a3"])
self.assertCountEqual(results, ["a1", "a2", "a3", "b1", "b2", "c1"])
async def test_send_starts_are_separated_by_configured_interval(self):
interval = 0.08
queue = await self._make_queue(interval_seconds=interval)
loop = asyncio.get_running_loop()
started_at: list[float] = []
async def operation() -> None:
started_at.append(loop.time())
tasks = [
asyncio.create_task(queue.submit(index + 1, operation, str(index)))
for index in range(3)
]
await asyncio.wait_for(asyncio.gather(*tasks), timeout=0.7)
self.assertEqual(len(started_at), 3)
gaps = [later - earlier for earlier, later in zip(started_at, started_at[1:])]
for gap in gaps:
# The Windows selector clock may wake one ~15.6 ms tick early.
self.assertGreaterEqual(gap, interval - 0.025)
async def test_operation_failure_does_not_block_later_job(self):
queue = await self._make_queue()
calls: list[str] = []
async def failing_operation() -> None:
calls.append("failed")
raise ValueError("expected failure")
async def following_operation() -> str:
calls.append("continued")
return "ok"
failed = asyncio.create_task(queue.submit(1, failing_operation, "failed"))
continued = asyncio.create_task(queue.submit(2, following_operation, "continued"))
results = await asyncio.wait_for(
asyncio.gather(failed, continued, return_exceptions=True),
timeout=0.2,
)
self.assertIsInstance(results[0], ValueError)
self.assertEqual(str(results[0]), "expected failure")
self.assertEqual(results[1], "ok")
self.assertEqual(calls, ["failed", "continued"])
async def test_cancelling_waiting_submit_prevents_operation(self):
queue = await self._make_queue()
first_started = asyncio.Event()
release_first = asyncio.Event()
cancelled_operation_called = asyncio.Event()
async def blocking_operation() -> None:
first_started.set()
await release_first.wait()
async def must_not_run() -> None:
cancelled_operation_called.set()
blocker = asyncio.create_task(queue.submit(1, blocking_operation, "blocker"))
await asyncio.wait_for(first_started.wait(), timeout=0.2)
waiting = asyncio.create_task(queue.submit(2, must_not_run, "cancelled"))
await self._wait_for_queued_count(queue, expected=1)
waiting.cancel()
with self.assertRaises(asyncio.CancelledError):
await waiting
release_first.set()
await asyncio.wait_for(blocker, timeout=0.2)
await asyncio.sleep(0.02)
self.assertFalse(cancelled_operation_called.is_set())
self.assertEqual((await queue.snapshot())["pending_count"], 0)
async def test_cancelling_active_submit_waits_for_real_result_without_overlap(self):
queue = await self._make_queue()
active = 0
maximum_active = 0
first_started = asyncio.Event()
release_first = asyncio.Event()
second_started = asyncio.Event()
async def first_operation() -> str:
nonlocal active, maximum_active
active += 1
maximum_active = max(maximum_active, active)
first_started.set()
await release_first.wait()
active -= 1
return "delivered"
async def second_operation() -> str:
nonlocal active, maximum_active
active += 1
maximum_active = max(maximum_active, active)
second_started.set()
await asyncio.sleep(0)
active -= 1
return "next"
first = asyncio.create_task(queue.submit(1, first_operation, "active"))
await asyncio.wait_for(first_started.wait(), timeout=0.2)
second = asyncio.create_task(queue.submit(2, second_operation, "next"))
await self._wait_for_queued_count(queue, expected=1)
first.cancel()
await asyncio.sleep(0.02)
# Once the network operation has begun, cancelling its caller must not
# release the single lane or report a false cancellation to the caller.
self.assertFalse(first.done())
self.assertFalse(second_started.is_set())
self.assertEqual(active, 1)
release_first.set()
self.assertEqual(await asyncio.wait_for(first, timeout=0.2), "delivered")
self.assertEqual(await asyncio.wait_for(second, timeout=0.2), "next")
self.assertTrue(second_started.is_set())
self.assertEqual(maximum_active, 1)
self.assertEqual(active, 0)
async def test_cancel_account_removes_waiting_jobs_without_running_them(self):
queue = await self._make_queue()
blocker_started = asyncio.Event()
release_blocker = asyncio.Event()
removed_calls: list[str] = []
async def blocking_operation() -> None:
blocker_started.set()
await release_blocker.wait()
async def removed_operation(name: str) -> None:
removed_calls.append(name)
blocker = asyncio.create_task(queue.submit(1, blocking_operation, "blocker"))
await asyncio.wait_for(blocker_started.wait(), timeout=0.2)
removed = [
asyncio.create_task(
queue.submit(
77,
lambda name=name: removed_operation(name),
name,
)
)
for name in ("waiting-1", "waiting-2")
]
await self._wait_for_queued_count(queue, expected=2)
self.assertEqual(await queue.cancel_account(77), 2)
results = await asyncio.wait_for(
asyncio.gather(*removed, return_exceptions=True),
timeout=0.1,
)
snapshot = await queue.snapshot()
self.assertTrue(all(isinstance(result, asyncio.CancelledError) for result in results))
self.assertEqual(removed_calls, [])
self.assertNotIn(77, snapshot["per_account"])
self.assertEqual(snapshot["queued_count"], 0)
self.assertEqual(snapshot["active_account_id"], 1)
release_blocker.set()
await asyncio.wait_for(blocker, timeout=0.2)
await asyncio.sleep(0)
self.assertEqual((await queue.snapshot())["pending_count"], 0)
async def test_cancel_account_releases_active_caller_but_drains_lane(self):
queue = await self._make_queue()
active = 0
maximum_active = 0
active_started = asyncio.Event()
release_active = asyncio.Event()
active_drained = asyncio.Event()
next_started = asyncio.Event()
calls: list[str] = []
async def active_operation() -> str:
nonlocal active, maximum_active
active += 1
maximum_active = max(maximum_active, active)
calls.append("active-start")
active_started.set()
await release_active.wait()
calls.append("active-drained")
active -= 1
active_drained.set()
return "discarded-result"
async def next_operation() -> str:
nonlocal active, maximum_active
active += 1
maximum_active = max(maximum_active, active)
calls.append("next")
next_started.set()
await asyncio.sleep(0)
active -= 1
return "next-result"
active_submit = asyncio.create_task(
queue.submit(88, active_operation, "active-account")
)
await asyncio.wait_for(active_started.wait(), timeout=0.2)
next_submit = asyncio.create_task(queue.submit(89, next_operation, "other-account"))
await self._wait_for_queued_count(queue, expected=1)
self.assertEqual(await queue.cancel_account(88), 1)
cancelled_result = await asyncio.wait_for(
asyncio.gather(active_submit, return_exceptions=True),
timeout=0.1,
)
self.assertIsInstance(cancelled_result[0], asyncio.CancelledError)
self.assertTrue(active_submit.cancelled())
self.assertFalse(active_drained.is_set())
self.assertFalse(next_started.is_set())
self.assertEqual(active, 1)
# The cancelled account's network operation still owns the lane until
# it actually drains; only then can a different account begin.
release_active.set()
self.assertEqual(
await asyncio.wait_for(next_submit, timeout=0.2),
"next-result",
)
self.assertTrue(active_drained.is_set())
self.assertEqual(calls, ["active-start", "active-drained", "next"])
self.assertEqual(maximum_active, 1)
self.assertEqual(active, 0)
# A fresh submission proves the dispatcher remains usable after the
# cancelled active job and its queued successor have both completed.
async def recovered_operation() -> str:
calls.append("recovered")
return "recovered-result"
self.assertEqual(
await asyncio.wait_for(
queue.submit(90, recovered_operation, "recovered"),
timeout=0.2,
),
"recovered-result",
)
self.assertEqual(calls[-1], "recovered")
self.assertEqual((await queue.snapshot())["pending_count"], 0)
class BackgroundTrafficLimitTests(unittest.IsolatedAsyncioTestCase):
async def _wait_until(self, predicate) -> None:
while not predicate():
await asyncio.sleep(0.001)
async def test_background_slot_respects_configured_concurrency_limit(self):
with patch.dict(
os.environ,
{"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "2"},
):
controller = TrafficController()
self.addAsyncCleanup(controller.stop)
active = 0
maximum_active = 0
entered = 0
all_finished = asyncio.Event()
async def request(account_id: int) -> None:
nonlocal active, maximum_active, entered
async with controller.background_slot(account_id, "poll"):
active += 1
entered += 1
maximum_active = max(maximum_active, active)
await asyncio.sleep(0.02)
active -= 1
if entered == 8 and active == 0:
all_finished.set()
tasks = [asyncio.create_task(request(index + 1)) for index in range(8)]
await asyncio.wait_for(asyncio.gather(*tasks), timeout=0.5)
self.assertEqual(entered, 8)
self.assertEqual(maximum_active, 2)
self.assertEqual(active, 0)
self.assertEqual(controller.background_active, 0)
self.assertEqual(controller.background_waiting, 0)
async def test_nested_background_slot_in_same_task_is_reentrant(self):
with patch.dict(
os.environ,
{"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "1"},
):
controller = TrafficController()
self.addAsyncCleanup(controller.stop)
tasks_seen: list[asyncio.Task] = []
async def nested_request() -> None:
tasks_seen.append(asyncio.current_task())
async with controller.background_slot(1, "outer"):
self.assertEqual(controller.background_active, 1)
self.assertEqual(controller.background_waiting, 0)
self.assertEqual(controller._background._value, 0)
tasks_seen.append(asyncio.current_task())
async with controller.background_slot(1, "inner"):
tasks_seen.append(asyncio.current_task())
self.assertEqual(controller.background_active, 1)
self.assertEqual(controller.background_waiting, 0)
self.assertEqual(controller._background._value, 0)
self.assertEqual(controller.background_active, 1)
self.assertEqual(controller._background._value, 0)
self.assertEqual(controller.background_active, 0)
self.assertEqual(controller.background_waiting, 0)
self.assertEqual(controller._background._value, 1)
await asyncio.wait_for(nested_request(), timeout=0.2)
self.assertEqual(len(tasks_seen), 3)
self.assertTrue(all(task is tasks_seen[0] for task in tasks_seen))
async def test_startup_traffic_leaves_one_slot_for_recurring_work(self):
with patch.dict(
os.environ,
{"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "3"},
):
controller = TrafficController()
self.addAsyncCleanup(controller.stop)
entered: list[str] = []
release = asyncio.Event()
async def request(name: str, *, startup: bool = False) -> None:
async with controller.background_slot(1, name, startup=startup):
entered.append(name)
await release.wait()
starts = [
asyncio.create_task(request(f"startup-{index}", startup=True))
for index in range(3)
]
while len(entered) < 2:
await asyncio.sleep(0)
await asyncio.sleep(0.01)
# A flood of startups may occupy at most capacity - 1 slots, so a
# recurring poll still gets in while the fleet is coming online.
self.assertEqual(len(entered), 2)
self.assertEqual(controller.background_startup_active, 2)
poll = asyncio.create_task(request("poll"))
await asyncio.wait_for(
self._wait_until(lambda: "poll" in entered),
timeout=0.5,
)
release.set()
await asyncio.wait_for(asyncio.gather(*starts, poll), timeout=0.5)
self.assertEqual(controller.background_active, 0)
self.assertEqual(controller.background_startup_active, 0)
async def test_recurring_work_is_not_deferred_indefinitely_by_startups(self):
"""Pending startup work may delay a recurring poll, never block it.
Starting several hundred accounts keeps startup requests queued for the
whole run. Yielding to that queue without a deadline left every hosted
account silent until the last account had finished coming online.
"""
with patch.dict(
os.environ,
{
"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "1",
"KEFU_BACKGROUND_NORMAL_MAX_DEFER_SECONDS": "0.02",
},
):
controller = TrafficController()
self.addAsyncCleanup(controller.stop)
entered: list[str] = []
release = asyncio.Event()
async def poll() -> None:
async with controller.background_slot(1, "conversation poll"):
entered.append("poll")
await release.wait()
# Stands in for a batch whose startup requests never stop arriving.
controller._background_startup_clear.clear()
task = asyncio.create_task(poll())
await asyncio.wait_for(
self._wait_until(lambda: entered == ["poll"]),
timeout=1.0,
)
release.set()
await asyncio.wait_for(task, timeout=0.5)
async def test_startup_request_is_not_buried_behind_normal_backlog(self):
with patch.dict(
os.environ,
{"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "1"},
):
controller = TrafficController()
self.addAsyncCleanup(controller.stop)
entered: list[str] = []
releases = {
name: asyncio.Event()
for name in ("active", "normal-1", "normal-2", "startup")
}
async def request(name: str, *, startup: bool = False) -> None:
async with controller.background_slot(
1,
name,
startup=startup,
):
entered.append(name)
await releases[name].wait()
active = asyncio.create_task(request("active"))
while entered != ["active"]:
await asyncio.sleep(0)
normal_one = asyncio.create_task(request("normal-1"))
normal_two = asyncio.create_task(request("normal-2"))
# Let one normal request reach the shared semaphore while the other is
# held at normal admission, then add the priority startup request.
await asyncio.sleep(0)
await asyncio.sleep(0)
startup = asyncio.create_task(request("startup", startup=True))
await asyncio.sleep(0)
releases["active"].set()
while len(entered) < 2:
await asyncio.sleep(0)
self.assertEqual(entered[:2], ["active", "normal-1"])
releases["normal-1"].set()
while len(entered) < 3:
await asyncio.sleep(0)
self.assertEqual(entered[:3], ["active", "normal-1", "startup"])
releases["startup"].set()
while len(entered) < 4:
await asyncio.sleep(0)
releases["normal-2"].set()
await asyncio.wait_for(
asyncio.gather(active, normal_one, normal_two, startup),
timeout=0.2,
)
self.assertEqual(
entered,
["active", "normal-1", "startup", "normal-2"],
)
self.assertEqual(controller.background_active, 0)
self.assertEqual(controller.background_waiting, 0)
self.assertEqual(controller.background_startup_active, 0)
self.assertEqual(controller.background_startup_waiting, 0)
class TrafficControllerLoopIsolationTests(unittest.TestCase):
def test_get_traffic_controller_does_not_reuse_asyncio_primitives(self):
async def capture_controller_state():
controller = get_traffic_controller()
return {
"loop": asyncio.get_running_loop(),
"controller": controller,
"send_lock": controller.send_queue._state_lock,
"send_wake": controller.send_queue._wake,
"background": controller._background,
"browser": controller._browser,
}
# IsolatedAsyncioTestCase uses a fresh asyncio.Runner per test. Keep
# both closed loop objects alive so the WeakKeyDictionary is exercised
# with two distinct keys rather than relying on garbage collection.
with asyncio.Runner() as first_runner:
first = first_runner.run(capture_controller_state())
with asyncio.Runner() as second_runner:
second = second_runner.run(capture_controller_state())
for key in (
"loop",
"controller",
"send_lock",
"send_wake",
"background",
"browser",
):
self.assertIsNot(first[key], second[key], key)
if __name__ == "__main__":
unittest.main()