319 lines
13 KiB
Python
319 lines
13 KiB
Python
"""Bounded, deduplicated background queue for bulk account starts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime, timezone
|
|
from typing import Any, Awaitable, Callable
|
|
|
|
|
|
logger = logging.getLogger("rpa.batch_start")
|
|
|
|
StartHandler = Callable[[int], Awaitable[dict[str, Any]]]
|
|
JobToken = tuple[str, int]
|
|
|
|
|
|
def _configured_concurrency() -> int:
|
|
try:
|
|
return max(1, min(8, int(os.getenv("KEFU_BATCH_START_CONCURRENCY", "2"))))
|
|
except (TypeError, ValueError):
|
|
return 2
|
|
|
|
|
|
def _utc_now() -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
@dataclass
|
|
class _BatchRecord:
|
|
batch_id: str
|
|
owner_id: int | None = None
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
items: dict[int, dict[str, Any]] = field(default_factory=dict)
|
|
created_at: str = field(default_factory=_utc_now)
|
|
updated_at: str = field(default_factory=_utc_now)
|
|
|
|
|
|
class BatchStartQueue:
|
|
"""Run account preparation with a small process-wide concurrency cap.
|
|
|
|
The HTTP endpoint only enqueues account ids. Long credential checks then
|
|
run in these workers, so a batch of many accounts cannot block the request
|
|
that submitted it. ``_pending_accounts`` makes overlapping clicks and
|
|
overlapping batches idempotent for each account. Pending and active
|
|
ownership is tied to a concrete job token so cleanup from a cancelled old
|
|
job cannot release a newer submission for the same account.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
handler: StartHandler,
|
|
concurrency: int | None = None,
|
|
max_batches: int = 100,
|
|
) -> None:
|
|
self._handler = handler
|
|
self.concurrency = max(1, int(concurrency or _configured_concurrency()))
|
|
self.max_batches = max(10, int(max_batches or 100))
|
|
self._queue: asyncio.Queue[JobToken] = asyncio.Queue()
|
|
self._pending_jobs: dict[int, JobToken] = {}
|
|
self._active_tasks: dict[int, tuple[JobToken, asyncio.Task]] = {}
|
|
self._batches: dict[str, _BatchRecord] = {}
|
|
self._workers: list[asyncio.Task] = []
|
|
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:
|
|
return
|
|
for index in range(self.concurrency):
|
|
self._workers.append(
|
|
asyncio.create_task(
|
|
self._worker(index + 1),
|
|
name=f"account-batch-start-{index + 1}",
|
|
)
|
|
)
|
|
|
|
async def submit(
|
|
self,
|
|
account_ids: list[int],
|
|
owner_id: int | None = None,
|
|
metadata: dict[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
await self._ensure_workers()
|
|
unique_ids = list(dict.fromkeys(int(value) for value in account_ids if int(value) > 0))
|
|
batch_id = uuid.uuid4().hex
|
|
record = _BatchRecord(
|
|
batch_id=batch_id,
|
|
owner_id=owner_id,
|
|
metadata=dict(metadata or {}),
|
|
)
|
|
|
|
async with self._lock:
|
|
if self._stopping:
|
|
raise RuntimeError("账号启动队列正在停止")
|
|
self._prune_locked()
|
|
self._batches[batch_id] = record
|
|
for account_id in unique_ids:
|
|
if account_id in self._pending_jobs:
|
|
record.items[account_id] = {
|
|
"account_id": account_id,
|
|
"status": "already_queued",
|
|
"message": "账号已在启动队列中",
|
|
}
|
|
continue
|
|
job_token = (batch_id, account_id)
|
|
self._pending_jobs[account_id] = job_token
|
|
record.items[account_id] = {
|
|
"account_id": account_id,
|
|
"status": "queued",
|
|
"message": "等待启动",
|
|
}
|
|
self._queue.put_nowait(job_token)
|
|
record.updated_at = _utc_now()
|
|
return self._snapshot_locked(record)
|
|
|
|
async def get_batch(
|
|
self,
|
|
batch_id: str,
|
|
owner_id: int | None = None,
|
|
include_items: bool = False,
|
|
) -> dict[str, Any] | None:
|
|
async with self._lock:
|
|
record = self._batches.get(str(batch_id or ""))
|
|
if record and owner_id is not None and record.owner_id != int(owner_id):
|
|
return None
|
|
return self._snapshot_locked(record, include_items=include_items) if record else None
|
|
|
|
async def _worker(self, worker_number: int) -> None:
|
|
while True:
|
|
batch_id, account_id = await self._queue.get()
|
|
job_token = (batch_id, account_id)
|
|
started_at = time.monotonic()
|
|
handler_task: asyncio.Task | None = None
|
|
try:
|
|
async with self._lock:
|
|
record = self._batches.get(batch_id)
|
|
if not record:
|
|
if self._pending_jobs.get(account_id) == job_token:
|
|
self._pending_jobs.pop(account_id, None)
|
|
continue
|
|
item = record.items[account_id]
|
|
if item.get("status") == "cancelled":
|
|
if self._pending_jobs.get(account_id) == job_token:
|
|
self._pending_jobs.pop(account_id, None)
|
|
continue
|
|
item.update(status="processing", message="正在校验并启动")
|
|
record.updated_at = _utc_now()
|
|
handler_task = asyncio.create_task(
|
|
self._handler(account_id),
|
|
name=f"account-start-{account_id}",
|
|
)
|
|
self._active_tasks[account_id] = (job_token, handler_task)
|
|
|
|
result = await handler_task
|
|
async with self._lock:
|
|
record = self._batches.get(batch_id)
|
|
if record:
|
|
item = record.items[account_id]
|
|
if item.get("status") != "cancelled":
|
|
item.update(
|
|
status="submitted",
|
|
message=str(result.get("message") or "已提交启动"),
|
|
login_mode=result.get("login_mode"),
|
|
skip_browser=bool(result.get("skip_browser", False)),
|
|
elapsed_seconds=round(time.monotonic() - started_at, 3),
|
|
)
|
|
record.updated_at = _utc_now()
|
|
except asyncio.CancelledError:
|
|
async with self._lock:
|
|
record = self._batches.get(batch_id)
|
|
if record:
|
|
record.items[account_id].update(
|
|
status="cancelled",
|
|
message="服务停止，启动任务已取消",
|
|
)
|
|
record.updated_at = _utc_now()
|
|
# Cancelling one account must not kill a long-lived queue
|
|
# worker. Re-raise only when stop() cancelled the worker.
|
|
if asyncio.current_task().cancelling():
|
|
raise
|
|
except Exception as exc:
|
|
logger.exception(
|
|
"Batch start failed account=%s worker=%s: %s",
|
|
account_id,
|
|
worker_number,
|
|
exc,
|
|
)
|
|
async with self._lock:
|
|
record = self._batches.get(batch_id)
|
|
if record:
|
|
detail = getattr(exc, "detail", None) or str(exc) or "启动失败"
|
|
item = record.items[account_id]
|
|
if item.get("status") != "cancelled":
|
|
item.update(
|
|
status="failed",
|
|
message=str(detail),
|
|
elapsed_seconds=round(time.monotonic() - started_at, 3),
|
|
)
|
|
record.updated_at = _utc_now()
|
|
finally:
|
|
async with self._lock:
|
|
active_entry = self._active_tasks.get(account_id)
|
|
if active_entry == (job_token, handler_task):
|
|
self._active_tasks.pop(account_id, None)
|
|
if self._pending_jobs.get(account_id) == job_token:
|
|
self._pending_jobs.pop(account_id, None)
|
|
self._queue.task_done()
|
|
|
|
async def cancel_account(self, account_id: int) -> int:
|
|
"""Cancel queued/processing work so stop/delete cannot restart it."""
|
|
account_key = int(account_id)
|
|
cancelled = 0
|
|
active_task: asyncio.Task | None = None
|
|
async with self._lock:
|
|
for record in self._batches.values():
|
|
item = record.items.get(account_key)
|
|
if not item or item.get("status") not in ("queued", "processing"):
|
|
continue
|
|
item.update(status="cancelled", message="启动任务已取消")
|
|
record.updated_at = _utc_now()
|
|
cancelled += 1
|
|
self._pending_jobs.pop(account_key, None)
|
|
active_entry = self._active_tasks.get(account_key)
|
|
active_task = active_entry[1] if active_entry else None
|
|
if active_task and not active_task.done():
|
|
active_task.cancel()
|
|
if active_task and not active_task.done():
|
|
await asyncio.gather(active_task, return_exceptions=True)
|
|
return cancelled
|
|
|
|
def _snapshot_locked(
|
|
self,
|
|
record: _BatchRecord,
|
|
*,
|
|
include_items: bool = True,
|
|
) -> dict[str, Any]:
|
|
counts = {
|
|
"queued": 0,
|
|
"processing": 0,
|
|
"submitted": 0,
|
|
"failed": 0,
|
|
"skipped": 0,
|
|
"cancelled": 0,
|
|
}
|
|
browser_required = 0
|
|
for item in record.items.values():
|
|
status = item.get("status")
|
|
if status == "submitted" and item.get("skip_browser") is False:
|
|
browser_required += 1
|
|
if status == "already_queued":
|
|
counts["skipped"] += 1
|
|
elif status in counts:
|
|
counts[status] += 1
|
|
active = counts["queued"] + counts["processing"]
|
|
snapshot = {
|
|
"batch_id": record.batch_id,
|
|
"total_count": len(record.items),
|
|
"accepted_count": len(record.items) - counts["skipped"],
|
|
"queued_count": counts["queued"],
|
|
"processing_count": counts["processing"],
|
|
"submitted_count": counts["submitted"],
|
|
"failed_count": counts["failed"],
|
|
"skipped_count": counts["skipped"],
|
|
"cancelled_count": counts["cancelled"],
|
|
"browser_required_count": browser_required,
|
|
"complete": active == 0,
|
|
"created_at": record.created_at,
|
|
"updated_at": record.updated_at,
|
|
}
|
|
snapshot.update(record.metadata)
|
|
if include_items:
|
|
snapshot["items"] = [dict(item) for item in record.items.values()]
|
|
return snapshot
|
|
|
|
def _prune_locked(self) -> None:
|
|
if len(self._batches) < self.max_batches:
|
|
return
|
|
removable = [
|
|
batch_id
|
|
for batch_id, record in self._batches.items()
|
|
if self._snapshot_locked(record, include_items=False)["complete"]
|
|
]
|
|
for batch_id in removable[: max(1, len(self._batches) - self.max_batches + 1)]:
|
|
self._batches.pop(batch_id, None)
|
|
|
|
async def stop(self) -> None:
|
|
async with self._lock:
|
|
self._stopping = True
|
|
workers = list(self._workers)
|
|
self._workers.clear()
|
|
for task in workers:
|
|
task.cancel()
|
|
if workers:
|
|
await asyncio.gather(*workers, return_exceptions=True)
|
|
async with self._lock:
|
|
self._active_tasks.clear()
|
|
while True:
|
|
try:
|
|
batch_id, account_id = self._queue.get_nowait()
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
record = self._batches.get(batch_id)
|
|
if record:
|
|
record.items[account_id].update(
|
|
status="cancelled",
|
|
message="服务停止，启动任务已取消",
|
|
)
|
|
record.updated_at = _utc_now()
|
|
job_token = (batch_id, account_id)
|
|
if self._pending_jobs.get(account_id) == job_token:
|
|
self._pending_jobs.pop(account_id, None)
|
|
self._queue.task_done()
|