diff --git a/backend/kefu.db-shm b/backend/kefu.db-shm new file mode 100644 index 0000000..3b97620 Binary files /dev/null and b/backend/kefu.db-shm differ diff --git a/backend/kefu.db-wal b/backend/kefu.db-wal new file mode 100644 index 0000000..36b8833 Binary files /dev/null and b/backend/kefu.db-wal differ diff --git a/backend/main.py b/backend/main.py index 566fa50..098bead 100644 --- a/backend/main.py +++ b/backend/main.py @@ -495,6 +495,22 @@ async def _conversations_from_logs( return list(seen.values()) +async def _release_db_connection(db: AsyncSession) -> None: + """Check this session's pooled connection back in before a long await. + + The whole backend shares a handful of pooled connections (five by default + on SQLite). A session left open across credential validation or worker + initialization holds one of them for tens of seconds per account, so every + unrelated request then waits out ``pool_timeout`` and the panel looks + frozen during a bulk start. Committing ends the transaction and returns + the connection; ``expire_on_commit=False`` keeps loaded attributes usable. + """ + try: + await db.commit() + except Exception: + logger.debug("Releasing the database connection failed", exc_info=True) + + async def _reset_account_credentials(account_id: int, db: AsyncSession) -> Account: """清除账号 Cookie、IM 会话等登录数据，并停止托管。""" await manager.stop_worker(account_id) @@ -1931,6 +1947,7 @@ async def update_account_cookie( # are committed so no worker can start in the stop/commit gap. async with manager.preparation_lock(account_id): account = await get_owned_account(db, user, account_id, write=True) + await _release_db_connection(db) await batch_start_queue.cancel_account(account_id) await manager.stop_worker(account_id) @@ -1970,6 +1987,7 @@ async def delete_account_cookie( # before the cleared credentials are committed. async with manager.preparation_lock(account_id): account = await get_owned_account(db, user, account_id, write=True) + await _release_db_connection(db) await batch_start_queue.cancel_account(account_id) await manager.stop_worker(account_id) @@ -2033,6 +2051,7 @@ async def delete_account( ): await get_owned_account(db, user, account_id, write=True) # 停止运行中的任务 + await _release_db_connection(db) await batch_start_queue.cancel_account(account_id) await manager.stop_worker(account_id) clear_cookie_file(account_id) @@ -2066,6 +2085,7 @@ async def reset_account_credentials( user: User = Depends(require_write), ): await get_owned_account(db, user, account_id, write=True) + await _release_db_connection(db) await batch_start_queue.cancel_account(account_id) account = await _reset_account_credentials(account_id, db) return { @@ -2087,6 +2107,9 @@ async def _start_account_rpa_impl( return {"status": "running", "message": "RPA worker is already running."} cookie_data = _get_account_cookie_data(account) + # Credential assessment issues real network requests to Douyin, and a bulk + # start runs it for every queued account. Release the connection first. + await _release_db_connection(db) assessment = await assess_account_credential( cookie_data, account.im_session_data, @@ -2136,6 +2159,10 @@ async def _start_account_rpa_impl( account.qr_code_base64 = None account.error_message = None await db.commit() + else: + # Nothing to persist, but the reads above may still hold a pooled + # connection, and waiting for readiness below takes seconds. + await _release_db_connection(db) try: started = await manager.start_worker( @@ -2213,6 +2240,10 @@ async def start_account_rpa( user: User = Depends(require_write), ): account = await get_owned_account(db, user, account_id, write=True) + # Cancelling waits for an in-flight queued start, and the preparation lock + # waits for whichever start owns this account. Neither may keep a pooled + # connection checked out while it waits. + await _release_db_connection(db) await batch_start_queue.cancel_account(account_id) async with manager.preparation_lock(account_id): return await _start_account_rpa_impl(account, db, body.login_mode) @@ -2290,6 +2321,9 @@ async def stop_account_rpa( ): account = await get_owned_account(db, user, account_id, write=True) + # Cancelling drains an in-flight queued start, which can take as long as + # the batch per-account deadline. Do not hold a pooled connection for it. + await _release_db_connection(db) await batch_start_queue.cancel_account(account_id) stopped = await manager.stop_worker(account_id) # 强制将数据库中的状态重置为 offline diff --git a/backend/rpa_engine/batch_start.py b/backend/rpa_engine/batch_start.py index 92c79b3..604e3df 100644 --- a/backend/rpa_engine/batch_start.py +++ b/backend/rpa_engine/batch_start.py @@ -22,11 +22,30 @@ class _StartPreparationTimeout(Exception): """Internal marker for the queue's own per-account deadline.""" +DEFAULT_CONCURRENCY = 6 +MAX_CONCURRENCY = 32 + + def _configured_concurrency() -> int: + """Admission width for account preparation. + + Every account spends most of its startup waiting: for the shared network + lane, for a WebSocket handshake, for signing work in a thread. Admitting + only two at a time therefore left the network lane idle and made a fleet of + several hundred accounts take tens of minutes. Actual outbound traffic is + still capped by the traffic controller, so a wider admission window fills + the existing lane instead of adding load. + """ try: - return max(1, min(8, int(os.getenv("KEFU_BATCH_START_CONCURRENCY", "2")))) + return max( + 1, + min( + MAX_CONCURRENCY, + int(os.getenv("KEFU_BATCH_START_CONCURRENCY", str(DEFAULT_CONCURRENCY))), + ), + ) except (TypeError, ValueError): - return 2 + return DEFAULT_CONCURRENCY def _configured_timeout_seconds() -> float: @@ -84,19 +103,25 @@ class BatchStartQueue: self._active_tasks: dict[int, tuple[JobToken, asyncio.Task]] = {} self._batches: dict[str, _BatchRecord] = {} self._workers: list[asyncio.Task] = [] + self._worker_sequence = 0 self._lock = asyncio.Lock() self._stopping = False async def _ensure_workers(self) -> None: async with self._lock: self._workers = [task for task in self._workers if not task.done()] - if self._workers or self._stopping: + if self._stopping: return - for index in range(self.concurrency): + # Top up to the configured width instead of only starting from + # zero. A worker that died on an unexpected error used to shrink + # the queue permanently, so later batches crawled through a single + # remaining worker with no way to recover short of a restart. + while len(self._workers) < self.concurrency: + self._worker_sequence += 1 self._workers.append( asyncio.create_task( - self._worker(index + 1), - name=f"account-batch-start-{index + 1}", + self._worker(self._worker_sequence), + name=f"account-batch-start-{self._worker_sequence}", ) ) diff --git a/backend/rpa_engine/douyin_im/traffic_control.py b/backend/rpa_engine/douyin_im/traffic_control.py index 15c9399..791bea9 100644 --- a/backend/rpa_engine/douyin_im/traffic_control.py +++ b/backend/rpa_engine/douyin_im/traffic_control.py @@ -393,9 +393,8 @@ class TrafficController: 1.0, ) ) - self._background = asyncio.Semaphore( - _env_int("KEFU_BACKGROUND_NETWORK_CONCURRENCY", 2) - ) + 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 @@ -405,6 +404,17 @@ class TrafficController: 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) ) @@ -449,8 +459,15 @@ class TrafficController: 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: @@ -460,7 +477,23 @@ class TrafficController: # 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: - await self._background_startup_clear.wait() + # 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 @@ -484,6 +517,7 @@ class TrafficController: self._background_owner.reset(token) if startup: self.background_startup_active -= 1 + self._background_startup.release() self.background_active -= 1 self._background.release() diff --git a/backend/tests/test_batch_start_api.py b/backend/tests/test_batch_start_api.py index d1b61de..40ad239 100644 --- a/backend/tests/test_batch_start_api.py +++ b/backend/tests/test_batch_start_api.py @@ -213,6 +213,55 @@ class BatchStartApiTests(unittest.IsolatedAsyncioTestCase): ) self.assertEqual(result["status"], "running") + async def test_batch_start_holds_no_db_connection_while_it_waits(self): + """A queued start must not pin one of the few pooled connections. + + Credential validation and readiness waiting take seconds per account. + Holding a session open across them exhausted the pool during a bulk + start, so every unrelated request waited out ``pool_timeout``. + """ + account = SimpleNamespace( + id=505, + status="starting", + qr_code_base64=None, + error_message=None, + im_session_data="saved-session", + ) + events: list[str] = [] + assessment = { + "login_mode": "im_direct", + "should_reset": False, + "can_skip_browser": True, + "message": "ready", + "cookie_valid": True, + "im_ready": True, + } + + async def commit(): + events.append("release") + + async def assess(*_args, **_kwargs): + events.append("assess") + return assessment + + async def start_worker(*_args, **_kwargs): + events.append("start-worker") + return True + + db = SimpleNamespace(commit=AsyncMock(side_effect=commit)) + + with ( + patch.object(main.manager, "is_running", return_value=False), + patch.object(main.manager, "start_worker", AsyncMock(side_effect=start_worker)), + patch.object(main, "_get_account_cookie_data", return_value="{}"), + patch.object(main, "assess_account_credential", AsyncMock(side_effect=assess)), + ): + await main._start_account_rpa_impl(account, db, wait_for_ready=True) + + # "starting" was already persisted, so the only commits here exist to + # return the connection: one before validation, one before the wait. + self.assertEqual(events, ["release", "assess", "release", "start-worker"]) + async def test_batch_start_does_not_launch_interactive_browser_login(self): account = SimpleNamespace( id=504, diff --git a/backend/tests/test_batch_start_queue.py b/backend/tests/test_batch_start_queue.py index 0043330..48ab781 100644 --- a/backend/tests/test_batch_start_queue.py +++ b/backend/tests/test_batch_start_queue.py @@ -87,6 +87,42 @@ class BatchStartQueueTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(maximum_active, 2) self.assertEqual(active, 0) + async def test_default_concurrency_admits_more_than_two_accounts(self): + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("KEFU_BATCH_START_CONCURRENCY", None) + queue = BatchStartQueue(lambda _account_id: None) + self.assertEqual(queue.concurrency, batch_start_module.DEFAULT_CONCURRENCY) + self.assertGreaterEqual(queue.concurrency, 4) + + async def test_dead_worker_is_replaced_so_width_never_shrinks(self): + """One crashed worker must not permanently narrow the queue. + + Width used to be restored only when every worker had exited, so a + single unexpected worker death left later batches crawling through the + survivors until the process restarted. + """ + + async def handler(account_id: int) -> dict: + return {"message": f"started-{account_id}"} + + queue = self._make_queue(handler, concurrency=3) + first = await queue.submit([91]) + await self._wait_for_complete(queue, first["batch_id"]) + self.assertEqual(len(queue._workers), 3) + + casualty = queue._workers[0] + casualty.cancel() + await asyncio.gather(casualty, return_exceptions=True) + + second = await queue.submit([92]) + await self._wait_for_complete(queue, second["batch_id"]) + + self.assertEqual(len(queue._workers), 3) + self.assertNotIn(casualty, queue._workers) + self.assertTrue(all(not task.done() for task in queue._workers)) + names = [task.get_name() for task in queue._workers] + self.assertEqual(len(set(names)), 3) + async def test_submit_returns_while_handler_is_blocked(self): handler_started = asyncio.Event() release_handler = asyncio.Event() diff --git a/backend/tests/test_cookie_credential_lock.py b/backend/tests/test_cookie_credential_lock.py index d4ebf7a..9645bb8 100644 --- a/backend/tests/test_cookie_credential_lock.py +++ b/backend/tests/test_cookie_credential_lock.py @@ -152,6 +152,9 @@ class CookieCredentialLockTests(unittest.IsolatedAsyncioTestCase): [ "lock-enter", "authorize", + # Checks the pooled connection back in before cancel/stop, + # which may wait on an in-flight start. + "commit", "cancel", "stop", "write-cookie", @@ -245,6 +248,9 @@ class CookieCredentialLockTests(unittest.IsolatedAsyncioTestCase): [ "lock-enter", "authorize", + # Checks the pooled connection back in before cancel/stop, + # which may wait on an in-flight start. + "commit", "cancel", "stop", "clear-cookie-file", diff --git a/backend/tests/test_traffic_control.py b/backend/tests/test_traffic_control.py index 0179f8f..3d5aef9 100644 --- a/backend/tests/test_traffic_control.py +++ b/backend/tests/test_traffic_control.py @@ -389,6 +389,10 @@ class GlobalSendQueueTests(unittest.IsolatedAsyncioTestCase): 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, @@ -458,6 +462,83 @@ class BackgroundTrafficLimitTests(unittest.IsolatedAsyncioTestCase): 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,