更新
This commit is contained in:
Your Name
@@ -0,0 +1,318 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user