feature: архитектура v2 — api + очередь + worker, удержание окна браузера (#78)
- Сервис api (FastAPI, без Chrome): POST /v2/tasks, GET /v2/tasks/{id} (доступ по cookie
сессии), DELETE /v2/sessions/current; /v1/solve — прежний контракт поверх очереди.
- Сервис worker (dramatiq + RabbitMQ): слоты браузера, окна сессий между задачами
(browser_idle_timeout 0–120 мин, простой — TTL ключа в Redis), heartbeat.
- Маршрутизация задач сессии в персональную очередь воркера с открытым окном.
- Redis: задачи, результаты (RESULT_TTL_HOURS), окна, воркеры. Прокси и ключ 2captcha
в Redis не сохраняются.
- Движок: AsyncStealthySession.run() — действия на текущей странице без навигации.
- Два образа (robotex-api, robotex-worker), compose с rabbitmq и redis,
make up/down/logs/smoke, релиз собирает все образы.
- Тесты: api + worker в одном процессе (fakeredis), настоящий воркер dramatiq на
StubBroker, браузерные тесты на новой схеме, smoke для запущенного стенда.
- Документация: README, API (v2), Deployment, Usage.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
1 parent
fe509c25c3
commit
7b511d5d38
38 files changed
+2142
-545
No files matched your search
@@ -0,0 +1,36 @@
|
||||
"""Запросы и ответы для лога без секретов клиента: креды прокси, ключ 2captcha, вводимые значения,
|
||||
значения cookies, html и скриншот."""
|
||||
import json
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from api.schemas import SolveRequest
|
||||
|
||||
|
||||
def mask_proxy(proxy_url: str | None) -> str | None:
|
||||
if not proxy_url:
|
||||
return proxy_url
|
||||
parts = urlsplit(proxy_url)
|
||||
return f"{parts.scheme}://{'***@' if parts.username else ''}{parts.hostname}:{parts.port}"
|
||||
|
||||
|
||||
def safe_request_dump(payload: SolveRequest) -> str:
|
||||
data = payload.model_dump(exclude_none=True)
|
||||
data["proxy"] = mask_proxy(payload.proxy)
|
||||
if payload.captcha:
|
||||
data["captcha"] = "***"
|
||||
if payload.cookies:
|
||||
data["cookies"] = [c.get("name") for c in payload.cookies]
|
||||
for act in data["actions"]:
|
||||
if act["type"] == "fill" and "value" in act:
|
||||
act["value"] = "***"
|
||||
return json.dumps(data, indent=2, ensure_ascii=False)
|
||||
|
||||
|
||||
def safe_response_dump(result: dict) -> str:
|
||||
data = dict(result)
|
||||
if data.get("cookies"):
|
||||
data["cookies"] = [c.get("name") for c in data["cookies"]]
|
||||
for key in ("html", "screenshot"):
|
||||
if data.get(key):
|
||||
data[key] = f"<{len(data[key])} chars>"
|
||||
return json.dumps(data, indent=2, ensure_ascii=False)
|
||||
+32
-1
@@ -1,9 +1,11 @@
|
||||
from typing import Any, Literal
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from pydantic import BaseModel, field_validator, model_validator
|
||||
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||
|
||||
CAPTCHA_TYPES = ("cloudflare", "recaptcha", "hcaptcha")
|
||||
# максимум browser_idle_timeout в минутах: открытое окно держит слот браузера (#78)
|
||||
MAX_BROWSER_IDLE_TIMEOUT = 120
|
||||
PROXY_SCHEMES = ("http", "https", "socks5")
|
||||
|
||||
# тип действия -> обязательные поля (кроме type)
|
||||
@@ -143,3 +145,32 @@ class SolveResponse(BaseModel):
|
||||
html: str | None = None
|
||||
screenshot: str | None = None
|
||||
duration_ms: int | None = None
|
||||
|
||||
|
||||
class TaskRequest(SolveRequest):
|
||||
"""
|
||||
Задача POST /v2/tasks: тело как у /v1/solve плюс удержание окна браузера.
|
||||
|
||||
url необязателен, если окно сессии открыто: тогда действия выполняются на текущей странице.
|
||||
"""
|
||||
url: str | None = None
|
||||
# минуты простоя, через которые окно сессии закроется; 0 — закрыть сразу после задачи
|
||||
browser_idle_timeout: int = Field(0, ge=0, le=MAX_BROWSER_IDLE_TIMEOUT)
|
||||
|
||||
|
||||
class TaskCreated(BaseModel):
|
||||
task_id: str
|
||||
status: str
|
||||
url: str
|
||||
|
||||
|
||||
class TaskResponse(BaseModel):
|
||||
"""
|
||||
GET /v2/tasks/{id}. status: queued → running → done | failed.
|
||||
done — результат в result (в нём свой status ok/err, как в /v1/solve);
|
||||
failed — задачу не удалось выполнить (воркер упал), причина в error.
|
||||
"""
|
||||
task_id: str
|
||||
status: Literal["queued", "running", "done", "failed", "cancelled"]
|
||||
error: str | None = None
|
||||
result: SolveResponse | None = None
|
||||
@@ -196,3 +196,23 @@ class AsyncStealthySession(AsyncSession, StealthySessionMixin):
|
||||
raise
|
||||
|
||||
raise RuntimeError("Request failed") # pragma: no cover
|
||||
|
||||
async def run(self, page_action, **kwargs: Unpack[StealthFetchParams]) -> Any:
|
||||
"""Выполнить `page_action` на текущей странице сессии, без навигации — для окна, удерживаемого
|
||||
между задачами (#82). Страница после `fetch()` остаётся в пуле открытой, её и берём.
|
||||
|
||||
:param page_action: функция от `page`; её результат возвращается.
|
||||
:param timeout / disable_resources / blocked_domains / extra_headers: как у `fetch()`.
|
||||
"""
|
||||
params = _validate(kwargs, self, StealthConfig)
|
||||
if not self._is_alive: # pragma: no cover
|
||||
raise RuntimeError("Context manager has been closed")
|
||||
|
||||
async with self._page_generator(
|
||||
params.timeout, params.extra_headers, params.disable_resources, None, params.blocked_domains
|
||||
) as page_info:
|
||||
try:
|
||||
return await page_action(page_info.page)
|
||||
except Exception:
|
||||
page_info.mark_error()
|
||||
raise
|
||||
+130
-200
@@ -1,237 +1,167 @@
|
||||
"""
|
||||
Сервис api (#78): регистрирует задачи в очереди RabbitMQ и отдаёт результаты из Redis. Chrome здесь
|
||||
нет — задачи выполняет сервис worker (src/worker).
|
||||
|
||||
* POST /v2/tasks, GET /v2/tasks/{id}, DELETE /v2/sessions/current — асинхронный API (#80, #82).
|
||||
* POST /v1/solve — прежний синхронный контракт поверх той же очереди (#84).
|
||||
"""
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from fastapi import Cookie
|
||||
from fastapi import FastAPI
|
||||
from fastapi import Response
|
||||
from fastapi import Cookie, FastAPI, HTTPException, Response
|
||||
from fastapi.responses import JSONResponse
|
||||
from patchright.async_api import Page, ProxySettings
|
||||
from starlette.concurrency import run_in_threadpool
|
||||
|
||||
logging.basicConfig(
|
||||
level=os.getenv("LOG_LEVEL", "INFO").upper(),
|
||||
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
|
||||
)
|
||||
|
||||
from api import actions, profile # noqa: E402 — модули ниже логируют при импорте, конфиг логов должен быть раньше
|
||||
from api.geoip import resolve_geo # noqa: E402
|
||||
from api.schemas import SolveRequest, SolveResponse # noqa: E402
|
||||
from engine import AsyncStealthySession # noqa: E402
|
||||
from api.schemas import SolveRequest, SolveResponse, TaskCreated, TaskRequest, TaskResponse # noqa: E402
|
||||
from tasks import broker as queue # noqa: E402
|
||||
from tasks.sessions import new_session_id, parse_session_id # noqa: E402
|
||||
from tasks.store import Status, TaskStore, connect # noqa: E402
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# сколько Chrome запускается параллельно; подбирать ≈ числу ядер (см. docs/Deployment.md)
|
||||
BROWSER_SLOTS = int(os.getenv("BROWSER_SLOTS", "10"))
|
||||
RETRY_AFTER = os.getenv("NO_SLOTS_RETRY_AFTER", "5")
|
||||
_slots = asyncio.Semaphore(BROWSER_SLOTS)
|
||||
|
||||
USER_DATA_ROOT = (Path(__file__).parent / "../extra/user_data_dir").resolve()
|
||||
|
||||
|
||||
def parse_proxy(proxy_url: str) -> ProxySettings:
|
||||
parts = urlsplit(proxy_url)
|
||||
server = f"{parts.scheme}://{parts.hostname}:{parts.port}"
|
||||
settings: ProxySettings = {"server": server}
|
||||
if parts.username:
|
||||
settings["username"] = parts.username
|
||||
if parts.password:
|
||||
settings["password"] = parts.password
|
||||
return settings
|
||||
|
||||
|
||||
def mask_proxy(proxy_url: str | None) -> str | None:
|
||||
if not proxy_url:
|
||||
return proxy_url
|
||||
parts = urlsplit(proxy_url)
|
||||
return f"{parts.scheme}://{'***@' if parts.username else ''}{parts.hostname}:{parts.port}"
|
||||
|
||||
|
||||
def safe_request_dump(payload: SolveRequest) -> str:
|
||||
"""Запрос для лога без секретов клиента: креды прокси, ключ 2captcha, вводимые значения, cookies."""
|
||||
data = payload.model_dump(exclude_none=True)
|
||||
data["proxy"] = mask_proxy(payload.proxy)
|
||||
if payload.captcha:
|
||||
data["captcha"] = "***"
|
||||
if payload.cookies:
|
||||
data["cookies"] = [c.get("name") for c in payload.cookies]
|
||||
for act in data["actions"]:
|
||||
if act["type"] == "fill" and "value" in act:
|
||||
act["value"] = "***"
|
||||
return json.dumps(data, indent=2, ensure_ascii=False)
|
||||
|
||||
|
||||
def safe_response_dump(result: dict) -> str:
|
||||
"""Ответ для лога: без значений cookies, html и скриншота."""
|
||||
data = dict(result)
|
||||
if data.get("cookies"):
|
||||
data["cookies"] = [c.get("name") for c in data["cookies"]]
|
||||
for key in ("html", "screenshot"):
|
||||
if data.get(key):
|
||||
data[key] = f"<{len(data[key])} chars>"
|
||||
return json.dumps(data, indent=2, ensure_ascii=False)
|
||||
|
||||
|
||||
def session_dir(session_id: str | None) -> tuple[str, Path]:
|
||||
"""Каталог профиля по session_id из cookie. Принимается только UUID — иначе (в т.ч. попытка
|
||||
подсунуть путь вида ../..) выдаётся новая сессия."""
|
||||
try:
|
||||
session_id = str(uuid.UUID(session_id)) if session_id else None
|
||||
except ValueError:
|
||||
logger.warning("invalid session_id cookie, issuing a new session")
|
||||
session_id = None
|
||||
session_id = session_id or str(uuid.uuid4())
|
||||
|
||||
user_data_dir = (USER_DATA_ROOT / session_id).resolve()
|
||||
if user_data_dir.parent != USER_DATA_ROOT:
|
||||
raise ValueError("invalid session path")
|
||||
return session_id, user_data_dir
|
||||
|
||||
|
||||
def make_page_action(payload: SolveRequest):
|
||||
"""Замыкание для AsyncStealthySession.fetch(page_action=...): выполняет очередь действий
|
||||
(включая прохождение капчи/антибота через api.actions.execute) и достаёт из уже
|
||||
навигированной страницы всё, что нужно для ответа API. Возврат этой функции — то, что
|
||||
вернёт fetch() (см. engine/stealthy.py).
|
||||
|
||||
Непройденный антибот не бросается наружу, а возвращается в challenge_error — чтобы в ответ
|
||||
всё равно попали cookies, html и скриншот страницы-заглушки."""
|
||||
|
||||
async def page_action(page: Page) -> dict:
|
||||
challenge_error = None
|
||||
try:
|
||||
actions_result = await actions.execute(
|
||||
page, payload.actions, payload.captcha,
|
||||
antibot_name=payload.captcha_type,
|
||||
success_locator=payload.captcha_success_locator,
|
||||
detect_locator=payload.captcha_detect_locator,
|
||||
)
|
||||
except actions.ChallengeFailed as e:
|
||||
challenge_error = str(e)
|
||||
actions_result = []
|
||||
|
||||
screenshot = None
|
||||
if payload.screenshot:
|
||||
screenshot = base64.b64encode(await page.screenshot(type="png")).decode()
|
||||
|
||||
return {
|
||||
"challenge_error": challenge_error,
|
||||
"actions": actions_result,
|
||||
"cookies": await page.context.cookies(),
|
||||
"html": await page.inner_html("html"),
|
||||
"screenshot": screenshot,
|
||||
"user_agent": await page.evaluate("() => navigator.userAgent"),
|
||||
}
|
||||
|
||||
return page_action
|
||||
# сколько /v1/solve ждёт, пока воркер возьмёт задачу, прежде чем ответить 503 no_slots
|
||||
V1_QUEUE_WAIT = float(os.getenv("V1_QUEUE_WAIT", "30"))
|
||||
POLL_INTERVAL = 0.2
|
||||
|
||||
store = TaskStore(connect())
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
NO_SLOTS = {"status": "err", "code": "no_slots", "error": "No slots available"}
|
||||
|
||||
|
||||
def no_slots_response() -> JSONResponse:
|
||||
return JSONResponse(status_code=503, headers={"Retry-After": RETRY_AFTER}, content=NO_SLOTS)
|
||||
|
||||
|
||||
def use_session(response: Response, cookie: str | None) -> str:
|
||||
"""session_id из cookie или новый; cookie ставится в ответ."""
|
||||
session_id = parse_session_id(cookie)
|
||||
if cookie and not session_id:
|
||||
logger.warning("invalid session_id cookie, issuing a new session")
|
||||
session_id = session_id or new_session_id()
|
||||
response.set_cookie("session_id", session_id)
|
||||
return session_id
|
||||
|
||||
|
||||
async def dispatch(message) -> None:
|
||||
# pika блокирующий — отправка в пуле потоков, чтобы не останавливать event loop
|
||||
await run_in_threadpool(queue.broker.enqueue, message)
|
||||
|
||||
|
||||
async def submit(payload: TaskRequest, session_id: str, window: dict | None) -> str:
|
||||
"""Зарегистрировать задачу и поставить её в очередь: в персональную очередь воркера, где открыто
|
||||
окно сессии (#83), иначе — в общую."""
|
||||
task_id = str(uuid.uuid4())
|
||||
await store.create(task_id, session_id, payload.url)
|
||||
target = queue.personal_queue(window["worker_id"]) if window else queue.COMMON_QUEUE
|
||||
await dispatch(queue.task_message(
|
||||
task_id, session_id, payload.model_dump(mode="json"), payload.browser_idle_timeout, target,
|
||||
))
|
||||
logger.info("task %s [%s] -> %s: %s", task_id, session_id, target, payload.url or "<current page>")
|
||||
return task_id
|
||||
|
||||
|
||||
async def free_slots() -> int:
|
||||
return sum(w["free_slots"] for w in await store.workers())
|
||||
|
||||
|
||||
@app.get("/")
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
return {"status": "ok", "browser_slots": BROWSER_SLOTS, "free_slots": _slots._value} # noqa
|
||||
workers = await store.workers()
|
||||
return {
|
||||
"status": "ok",
|
||||
"browser_slots": sum(w["browser_slots"] for w in workers),
|
||||
"free_slots": sum(w["free_slots"] for w in workers),
|
||||
"workers": workers,
|
||||
}
|
||||
|
||||
|
||||
@app.post("/v2/tasks", status_code=202, response_model=TaskCreated)
|
||||
async def create_task(
|
||||
payload: TaskRequest,
|
||||
response: Response,
|
||||
session_id: Annotated[str | None, Cookie()] = None,
|
||||
):
|
||||
# TODO: добавить поддержку api ключей (#4)
|
||||
session_id = use_session(response, session_id)
|
||||
window = await store.live_window(session_id)
|
||||
if payload.url is None and window is None:
|
||||
raise HTTPException(422, "url is required: the session has no open browser window")
|
||||
task_id = await submit(payload, session_id, window)
|
||||
return {"task_id": task_id, "status": Status.QUEUED, "url": f"/v2/tasks/{task_id}"}
|
||||
|
||||
|
||||
@app.get("/v2/tasks/{task_id}", response_model=TaskResponse, response_model_exclude_none=True)
|
||||
async def get_task(task_id: str, session_id: Annotated[str | None, Cookie()] = None):
|
||||
# задача видна только своей сессии; чужая и несуществующая неразличимы
|
||||
task = await store.get(task_id) if parse_session_id(task_id) else None
|
||||
if not task or task["session_id"] != parse_session_id(session_id):
|
||||
raise HTTPException(404, "task not found")
|
||||
return {"task_id": task_id, "status": task["status"], "error": task.get("error"), "result": task.get("result")}
|
||||
|
||||
|
||||
@app.delete("/v2/sessions/current")
|
||||
async def close_session(session_id: Annotated[str | None, Cookie()] = None):
|
||||
"""Закрыть окно сессии досрочно и освободить слот."""
|
||||
session_id = parse_session_id(session_id)
|
||||
window = await store.get_session(session_id) if session_id else None
|
||||
if not window:
|
||||
return {"closed": False}
|
||||
await store.delete_session(session_id)
|
||||
if await store.worker_alive(window["worker_id"]):
|
||||
await dispatch(queue.close_message(session_id, queue.personal_queue(window["worker_id"])))
|
||||
return {"closed": True}
|
||||
|
||||
|
||||
@app.post("/v1/solve", response_model=SolveResponse, response_model_exclude_none=True)
|
||||
async def solve(
|
||||
payload: SolveRequest,
|
||||
response: Response,
|
||||
session_id: Annotated[str | None, Cookie()] = None
|
||||
session_id: Annotated[str | None, Cookie()] = None,
|
||||
):
|
||||
"""Синхронный контракт v1: задача ставится в ту же очередь, ответ — когда воркер её выполнит.
|
||||
Если свободных слотов нет или воркер не взял задачу за V1_QUEUE_WAIT — 503, как и раньше."""
|
||||
# TODO: добавить поддержку api ключей (#4)
|
||||
started = time.monotonic()
|
||||
session_id = use_session(response, session_id)
|
||||
window = await store.live_window(session_id)
|
||||
if window is None and await free_slots() <= 0:
|
||||
logger.warning("no free browser slots, rejecting %s", payload.url)
|
||||
return no_slots_response()
|
||||
|
||||
if _slots.locked():
|
||||
logger.warning("no free browser slots (%d), rejecting %s", BROWSER_SLOTS, payload.url)
|
||||
return JSONResponse(
|
||||
status_code=503,
|
||||
headers={"Retry-After": RETRY_AFTER},
|
||||
content={"status": "err", "code": "no_slots", "error": "No slots available"},
|
||||
)
|
||||
task_id = await submit(TaskRequest.model_validate(payload.model_dump()), session_id, window)
|
||||
while True:
|
||||
task = await store.get(task_id)
|
||||
if task is None:
|
||||
result = {"status": "err", "code": "request_failed", "error": "task expired"}
|
||||
break
|
||||
if task["status"] == Status.DONE:
|
||||
result = task["result"]
|
||||
break
|
||||
if task["status"] in (Status.FAILED, Status.CANCELLED):
|
||||
result = {"status": "err", "code": "request_failed", "error": task.get("error") or task["status"]}
|
||||
break
|
||||
if (
|
||||
task["status"] == Status.QUEUED
|
||||
and time.monotonic() - started > V1_QUEUE_WAIT
|
||||
and await store.cancel(task_id)
|
||||
):
|
||||
logger.warning("task %s was not picked up in %ss, rejecting", task_id, V1_QUEUE_WAIT)
|
||||
return no_slots_response()
|
||||
await asyncio.sleep(POLL_INTERVAL)
|
||||
|
||||
async with _slots:
|
||||
session_id, user_data_dir = session_dir(session_id)
|
||||
response.set_cookie("session_id", session_id)
|
||||
logger.info("Incoming request [%s]: %s", session_id, payload.url)
|
||||
logger.debug("Incoming payload [%s]:\n%s", session_id, safe_request_dump(payload))
|
||||
|
||||
try:
|
||||
user_data_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# отпечаток задаётся профилю один раз (см. api/profile.py)
|
||||
profile_kwargs = profile.browser_kwargs(profile.load_or_create(user_data_dir))
|
||||
|
||||
proxy = parse_proxy(payload.proxy) if payload.proxy else None
|
||||
|
||||
# geo — по exit-IP прокси (см. api/geoip.py); без прокси оставляем настройки браузера
|
||||
geo = await resolve_geo(proxy) if proxy else None
|
||||
locale = geo.get("locale") if geo else None
|
||||
|
||||
nav_timeout_ms = payload.timeout * 1000 if payload.timeout else 60_000
|
||||
blocked_domains = set(payload.blocked_domains) if payload.blocked_domains else None
|
||||
|
||||
kwargs = dict(
|
||||
headless=False,
|
||||
proxy=proxy,
|
||||
user_data_dir=str(user_data_dir),
|
||||
real_chrome=True,
|
||||
locale=locale,
|
||||
timezone_id=geo.get("timezone_id") if geo else None,
|
||||
cookies=payload.cookies,
|
||||
# переданные cookies заменяют cookies профиля, а не дополняют их
|
||||
clear_cookies=payload.cookies is not None,
|
||||
disable_resources=payload.disable_resources,
|
||||
blocked_domains=blocked_domains,
|
||||
timeout=nav_timeout_ms,
|
||||
retries=max(1, payload.retries),
|
||||
retry_delay=payload.retry_delay,
|
||||
wait=payload.wait,
|
||||
google_search=payload.google_search,
|
||||
**profile_kwargs,
|
||||
)
|
||||
|
||||
async with AsyncStealthySession(**kwargs) as engine:
|
||||
result = await engine.fetch(payload.url, page_action=make_page_action(payload))
|
||||
except Exception as e:
|
||||
logger.exception("request failed [%s]", session_id)
|
||||
result_payload = {
|
||||
"status": "err",
|
||||
"code": "request_failed",
|
||||
"error": str(e) or type(e).__name__,
|
||||
}
|
||||
else:
|
||||
failed = [a for a in result["actions"] if a.get("ok") is False]
|
||||
result_payload = {
|
||||
"status": "ok",
|
||||
"user_agent": result["user_agent"],
|
||||
"cookies": result["cookies"],
|
||||
"actions": result["actions"],
|
||||
}
|
||||
if result["challenge_error"]:
|
||||
result_payload.update(status="err", code="challenge_failed", error=result["challenge_error"])
|
||||
elif failed:
|
||||
result_payload.update(
|
||||
status="err", code="action_failed",
|
||||
error=f"{failed[0]['type']}: {failed[0].get('result')}",
|
||||
)
|
||||
if payload.html:
|
||||
result_payload["html"] = result["html"]
|
||||
if payload.screenshot:
|
||||
result_payload["screenshot"] = result["screenshot"]
|
||||
|
||||
result_payload["duration_ms"] = int((time.monotonic() - started) * 1000)
|
||||
logger.info(
|
||||
"Outgoing [%s]: %s %s %dms", session_id, result_payload["status"],
|
||||
result_payload.get("error", ""), result_payload["duration_ms"],
|
||||
)
|
||||
logger.debug("Outgoing payload:\n%s", safe_response_dump(result_payload))
|
||||
return result_payload
|
||||
result.setdefault("actions", [])
|
||||
result["duration_ms"] = int((time.monotonic() - started) * 1000)
|
||||
return result
|
||||
@@ -0,0 +1,5 @@
|
||||
"""
|
||||
Общее для сервисов api и worker (#78): хранилище задач в Redis, брокер dramatiq, сессии.
|
||||
|
||||
Пакет не зависит от браузера: образ api собирается без Chrome и Patchright.
|
||||
"""
|
||||
@@ -0,0 +1,52 @@
|
||||
"""
|
||||
Брокер dramatiq поверх RabbitMQ (#78) и сообщения, которыми api ставит работу воркерам.
|
||||
|
||||
Каждый воркер слушает общую очередь `robotex` и персональную `robotex.worker.<id>`. В персональную
|
||||
api отправляет задачи сессии, чьё окно открыто на этом воркере (#83), иначе — в общую. Сообщение
|
||||
адресуется actor'у по имени, поэтому api не импортирует код воркера (и Chrome ему не нужен).
|
||||
|
||||
`RABBITMQ_URL=stub://` — брокер в памяти для тестов.
|
||||
"""
|
||||
import os
|
||||
|
||||
import dramatiq
|
||||
from dramatiq.brokers.stub import StubBroker
|
||||
from dramatiq.middleware import AsyncIO
|
||||
|
||||
COMMON_QUEUE = "robotex"
|
||||
RUN_TASK = "run_task"
|
||||
CLOSE_WINDOW = "close_window"
|
||||
|
||||
|
||||
def personal_queue(worker_id: str) -> str:
|
||||
return f"{COMMON_QUEUE}.worker.{worker_id}"
|
||||
|
||||
|
||||
def make_broker(url: str) -> dramatiq.Broker:
|
||||
if url.startswith("stub://"):
|
||||
broker = StubBroker()
|
||||
else:
|
||||
from dramatiq.brokers.rabbitmq import RabbitmqBroker
|
||||
# confirm_delivery: enqueue возвращается, только когда RabbitMQ принял сообщение
|
||||
broker = RabbitmqBroker(url=url, confirm_delivery=True)
|
||||
broker.add_middleware(AsyncIO())
|
||||
return broker
|
||||
|
||||
|
||||
broker = make_broker(os.getenv("RABBITMQ_URL", "amqp://guest:guest@localhost:5672/"))
|
||||
dramatiq.set_broker(broker)
|
||||
|
||||
|
||||
def task_message(
|
||||
task_id: str, session_id: str, payload: dict, idle_timeout: int, queue: str = COMMON_QUEUE,
|
||||
) -> dramatiq.Message:
|
||||
"""payload — SolveRequest/TaskRequest.model_dump(): в нём прокси и ключ 2captcha, поэтому
|
||||
он живёт только в сообщении, а не в Redis."""
|
||||
return dramatiq.Message(
|
||||
queue_name=queue, actor_name=RUN_TASK,
|
||||
args=(task_id, session_id, payload, idle_timeout), kwargs={}, options={},
|
||||
)
|
||||
|
||||
|
||||
def close_message(session_id: str, queue: str) -> dramatiq.Message:
|
||||
return dramatiq.Message(queue_name=queue, actor_name=CLOSE_WINDOW, args=(session_id,), kwargs={}, options={})
|
||||
@@ -0,0 +1,16 @@
|
||||
import uuid
|
||||
|
||||
|
||||
def parse_session_id(value: str | None) -> str | None:
|
||||
"""session_id из cookie. Принимается только UUID: id становится именем каталога профиля, и строка
|
||||
вида ../.. не должна до него дойти. Невалидное значение — как отсутствие сессии."""
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
return str(uuid.UUID(value))
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def new_session_id() -> str:
|
||||
return str(uuid.uuid4())
|
||||
@@ -0,0 +1,199 @@
|
||||
"""
|
||||
Состояние задач, окон и воркеров в Redis (#78). Сроки жизни — TTL ключей Redis.
|
||||
|
||||
* `robotex:task:<id>` (hash) — статус задачи, сессия, время, результат. Пока задача не завершена,
|
||||
у ключа страховочный TTL `PENDING_TTL`; при завершении — `RESULT_TTL` (`RESULT_TTL_HOURS`).
|
||||
* `robotex:session:<id>` (hash) — где открыто окно сессии: `worker_id` и `idle_timeout`.
|
||||
Пока задача выполняется, TTL — `RUNNING_TTL`; после неё — `browser_idle_timeout`. Ключ исчез —
|
||||
окно закрывается (worker/windows.py).
|
||||
* `robotex:worker:<id>` — heartbeat воркера с TTL `HEARTBEAT_TTL`: слоты и число окон.
|
||||
|
||||
Прокси и ключ 2captcha сюда не попадают: они передаются только в сообщении очереди.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
|
||||
from redis.asyncio import Redis
|
||||
from redis.exceptions import WatchError
|
||||
|
||||
PREFIX = "robotex"
|
||||
|
||||
RESULT_TTL = int(float(os.getenv("RESULT_TTL_HOURS", "24")) * 3600)
|
||||
# задача не может висеть в очереди или выполняться дольше: дальше её ключ просто исчезает
|
||||
PENDING_TTL = 6 * 3600
|
||||
# пока задача выполняется, ключ сессии держит маршрут к воркеру; дольше задача не живёт (TimeLimit)
|
||||
RUNNING_TTL = 3 * 3600
|
||||
HEARTBEAT_TTL = 30
|
||||
# сколько раз задачу можно начать заново после падения воркера, прежде чем считать её проваленной
|
||||
MAX_ATTEMPTS = 2
|
||||
|
||||
|
||||
class Status:
|
||||
QUEUED = "queued"
|
||||
RUNNING = "running"
|
||||
DONE = "done"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
def connect(url: str | None = None) -> Redis:
|
||||
return Redis.from_url(url or os.getenv("REDIS_URL", "redis://localhost:6379/0"), decode_responses=True)
|
||||
|
||||
|
||||
def _task_key(task_id: str) -> str:
|
||||
return f"{PREFIX}:task:{task_id}"
|
||||
|
||||
|
||||
def _session_key(session_id: str) -> str:
|
||||
return f"{PREFIX}:session:{session_id}"
|
||||
|
||||
|
||||
def _worker_key(worker_id: str) -> str:
|
||||
return f"{PREFIX}:worker:{worker_id}"
|
||||
|
||||
|
||||
class TaskStore:
|
||||
def __init__(self, redis: Redis):
|
||||
self.redis = redis
|
||||
|
||||
async def ping(self) -> bool:
|
||||
return await self.redis.ping()
|
||||
|
||||
# --- задачи ---
|
||||
|
||||
async def create(self, task_id: str, session_id: str, url: str | None) -> None:
|
||||
key = _task_key(task_id)
|
||||
async with self.redis.pipeline(transaction=True) as pipe:
|
||||
pipe.hset(key, mapping={
|
||||
"status": Status.QUEUED,
|
||||
"session_id": session_id,
|
||||
"url": url or "",
|
||||
"created_at": time.time(),
|
||||
"attempts": 0,
|
||||
})
|
||||
pipe.expire(key, PENDING_TTL)
|
||||
await pipe.execute()
|
||||
|
||||
async def get(self, task_id: str) -> dict | None:
|
||||
data = await self.redis.hgetall(_task_key(task_id))
|
||||
if not data:
|
||||
return None
|
||||
if "result" in data:
|
||||
data["result"] = json.loads(data["result"])
|
||||
return data
|
||||
|
||||
async def claim(self, task_id: str, worker_id: str, running_here: bool = False) -> bool:
|
||||
"""Взять задачу в работу: queued → running. Атомарно (WATCH), чтобы задачу, отменённую
|
||||
/v1/solve по таймауту очереди, никто не начал выполнять.
|
||||
|
||||
Задачу в статусе running можно взять заново, если её воркер умер (RabbitMQ вернул
|
||||
неподтверждённое сообщение в очередь) — не больше MAX_ATTEMPTS раз. running_here — задача
|
||||
уже выполняется в этом процессе: повторная доставка того же сообщения, второй раз не берём."""
|
||||
key = _task_key(task_id)
|
||||
async with self.redis.pipeline(transaction=True) as pipe:
|
||||
while True:
|
||||
try:
|
||||
await pipe.watch(key)
|
||||
data = await pipe.hgetall(key)
|
||||
if not data:
|
||||
return False
|
||||
status = data["status"]
|
||||
if status == Status.RUNNING:
|
||||
owner = data.get("worker_id")
|
||||
if running_here or (owner != worker_id and await self.worker_alive(owner)):
|
||||
return False
|
||||
if int(data.get("attempts", 0)) >= MAX_ATTEMPTS:
|
||||
pipe.multi()
|
||||
self._finish(pipe, key, Status.FAILED, error="worker died while running the task")
|
||||
await pipe.execute()
|
||||
return False
|
||||
elif status != Status.QUEUED:
|
||||
return False
|
||||
pipe.multi()
|
||||
pipe.hset(key, mapping={"status": Status.RUNNING, "worker_id": worker_id, "started_at": time.time()})
|
||||
pipe.hincrby(key, "attempts", 1)
|
||||
await pipe.execute()
|
||||
return True
|
||||
except WatchError:
|
||||
continue
|
||||
|
||||
async def cancel(self, task_id: str) -> bool:
|
||||
"""queued → cancelled; False, если воркер уже взял задачу."""
|
||||
key = _task_key(task_id)
|
||||
async with self.redis.pipeline(transaction=True) as pipe:
|
||||
while True:
|
||||
try:
|
||||
await pipe.watch(key)
|
||||
if await pipe.hget(key, "status") != Status.QUEUED:
|
||||
return False
|
||||
pipe.multi()
|
||||
self._finish(pipe, key, Status.CANCELLED)
|
||||
await pipe.execute()
|
||||
return True
|
||||
except WatchError:
|
||||
continue
|
||||
|
||||
async def finish(self, task_id: str, result: dict) -> None:
|
||||
async with self.redis.pipeline(transaction=True) as pipe:
|
||||
self._finish(pipe, _task_key(task_id), Status.DONE, result=result)
|
||||
await pipe.execute()
|
||||
|
||||
async def fail(self, task_id: str, error: str) -> None:
|
||||
async with self.redis.pipeline(transaction=True) as pipe:
|
||||
self._finish(pipe, _task_key(task_id), Status.FAILED, error=error)
|
||||
await pipe.execute()
|
||||
|
||||
@staticmethod
|
||||
def _finish(pipe, key: str, status: str, result: dict | None = None, error: str | None = None) -> None:
|
||||
fields = {"status": status, "finished_at": time.time()}
|
||||
if result is not None:
|
||||
fields["result"] = json.dumps(result, ensure_ascii=False)
|
||||
if error is not None:
|
||||
fields["error"] = error
|
||||
pipe.hset(key, mapping=fields)
|
||||
pipe.expire(key, RESULT_TTL)
|
||||
|
||||
# --- окна сессий ---
|
||||
|
||||
async def set_session(self, session_id: str, worker_id: str, idle_timeout: int, ttl: int) -> None:
|
||||
key = _session_key(session_id)
|
||||
async with self.redis.pipeline(transaction=True) as pipe:
|
||||
pipe.hset(key, mapping={"worker_id": worker_id, "idle_timeout": idle_timeout})
|
||||
pipe.expire(key, ttl)
|
||||
await pipe.execute()
|
||||
|
||||
async def get_session(self, session_id: str) -> dict | None:
|
||||
return await self.redis.hgetall(_session_key(session_id)) or None
|
||||
|
||||
async def delete_session(self, session_id: str) -> None:
|
||||
await self.redis.delete(_session_key(session_id))
|
||||
|
||||
async def session_ttl(self, session_id: str) -> int:
|
||||
return await self.redis.ttl(_session_key(session_id))
|
||||
|
||||
async def live_window(self, session_id: str) -> dict | None:
|
||||
"""Окно сессии, если оно открыто на живом воркере: {"worker_id", "idle_timeout"}."""
|
||||
window = await self.get_session(session_id)
|
||||
if window and await self.worker_alive(window["worker_id"]):
|
||||
return window
|
||||
return None
|
||||
|
||||
# --- воркеры ---
|
||||
|
||||
async def heartbeat(self, worker_id: str, info: dict) -> None:
|
||||
await self.redis.set(_worker_key(worker_id), json.dumps(info), ex=HEARTBEAT_TTL)
|
||||
|
||||
async def remove_worker(self, worker_id: str) -> None:
|
||||
await self.redis.delete(_worker_key(worker_id))
|
||||
|
||||
async def worker_alive(self, worker_id: str | None) -> bool:
|
||||
return bool(worker_id) and bool(await self.redis.exists(_worker_key(worker_id)))
|
||||
|
||||
async def workers(self) -> list[dict]:
|
||||
result = []
|
||||
async for key in self.redis.scan_iter(match=_worker_key("*")):
|
||||
info = await self.redis.get(key)
|
||||
if info:
|
||||
result.append({"worker_id": key.rsplit(":", 1)[-1], **json.loads(info)})
|
||||
return sorted(result, key=lambda w: w["worker_id"])
|
||||
@@ -0,0 +1,6 @@
|
||||
"""
|
||||
Сервис worker (#78): берёт задачи из очереди RabbitMQ, выполняет их в Chrome и кладёт результат в Redis.
|
||||
|
||||
Запуск: `dramatiq worker.actors --processes 1 --threads N` под Xvfb (см. Dockerfile). Процесс — строго
|
||||
один: открытые окна браузера живут в его памяти.
|
||||
"""
|
||||
@@ -0,0 +1,94 @@
|
||||
"""
|
||||
Actor'ы dramatiq воркера. Точка входа: `dramatiq worker.actors --processes 1 --threads N`.
|
||||
|
||||
Actor'ы синхронные: каждый передаёт работу в общий event loop (поток AsyncIO-middleware), где живут
|
||||
окна браузера, и ждёт её. Так возврат задачи в очередь (enqueue по pika, блокирующий) идёт из потока
|
||||
dramatiq, а не из event loop.
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
|
||||
logging.basicConfig(
|
||||
level=os.getenv("LOG_LEVEL", "INFO").upper(),
|
||||
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
|
||||
)
|
||||
|
||||
import dramatiq # noqa: E402
|
||||
from dramatiq.asyncio import get_event_loop_thread # noqa: E402
|
||||
from dramatiq.middleware import Middleware # noqa: E402
|
||||
|
||||
from tasks import store as task_store # noqa: E402
|
||||
from tasks.broker import CLOSE_WINDOW, COMMON_QUEUE, RUN_TASK, broker, personal_queue # noqa: E402
|
||||
from worker.windows import Outcome, WindowPool # noqa: E402
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
WORKER_ID = os.getenv("WORKER_ID") or socket.gethostname()
|
||||
# сколько окон Chrome одновременно — и выполняющих задачу, и удерживаемых; ≈ числу ядер (docs/Deployment.md)
|
||||
BROWSER_SLOTS = int(os.getenv("BROWSER_SLOTS", "4"))
|
||||
# через сколько задача без свободного слота снова попадёт в очередь
|
||||
NO_SLOT_DELAY_MS = 2000
|
||||
HEARTBEAT_INTERVAL = 10
|
||||
TASK_TIME_LIMIT_MS = task_store.RUNNING_TTL * 1000
|
||||
|
||||
_pool: WindowPool | None = None
|
||||
_background: asyncio.Task | None = None
|
||||
|
||||
|
||||
def _run(coro):
|
||||
return get_event_loop_thread().run_coroutine(coro)
|
||||
|
||||
|
||||
async def _heartbeat_loop(pool: WindowPool) -> None:
|
||||
"""heartbeat воркера (по нему api видит, жив ли воркер с окном) и закрытие простаивающих окон."""
|
||||
while True:
|
||||
try:
|
||||
await pool.report()
|
||||
await pool.expire_idle()
|
||||
except Exception:
|
||||
log.warning("idle check failed", exc_info=True)
|
||||
await asyncio.sleep(HEARTBEAT_INTERVAL)
|
||||
|
||||
|
||||
async def _start() -> None:
|
||||
global _pool, _background
|
||||
_pool = WindowPool(task_store.TaskStore(task_store.connect()), WORKER_ID, BROWSER_SLOTS)
|
||||
await _pool.report()
|
||||
_background = asyncio.create_task(_heartbeat_loop(_pool))
|
||||
log.info("worker %s: %d browser slots, queues %s, %s", WORKER_ID, BROWSER_SLOTS, COMMON_QUEUE, personal_queue(WORKER_ID))
|
||||
|
||||
|
||||
async def _stop() -> None:
|
||||
if _background:
|
||||
_background.cancel()
|
||||
await _pool.close_all()
|
||||
await _pool.store.remove_worker(WORKER_ID)
|
||||
|
||||
|
||||
class Lifecycle(Middleware):
|
||||
"""Пул окон и heartbeat — в event loop AsyncIO-middleware: он поднимается раньше (middleware
|
||||
добавлен в брокер раньше), а окна закрываются до его остановки."""
|
||||
|
||||
def before_worker_boot(self, broker, worker):
|
||||
_run(_start())
|
||||
|
||||
def before_worker_shutdown(self, broker, worker):
|
||||
_run(_stop())
|
||||
|
||||
|
||||
broker.add_middleware(Lifecycle())
|
||||
broker.declare_queue(personal_queue(WORKER_ID))
|
||||
|
||||
|
||||
@dramatiq.actor(actor_name=RUN_TASK, queue_name=COMMON_QUEUE, max_retries=0, time_limit=TASK_TIME_LIMIT_MS)
|
||||
def run_task(task_id: str, session_id: str, payload: dict, idle_timeout: int) -> None:
|
||||
if _run(_pool.handle(task_id, session_id, payload, idle_timeout)) == Outcome.NO_SLOT:
|
||||
# в общую очередь: окно могло открыться на любом воркере, где есть слот
|
||||
run_task.send_with_options(args=(task_id, session_id, payload, idle_timeout), delay=NO_SLOT_DELAY_MS)
|
||||
|
||||
|
||||
@dramatiq.actor(actor_name=CLOSE_WINDOW, queue_name=COMMON_QUEUE, max_retries=0)
|
||||
def close_window(session_id: str) -> None:
|
||||
_run(_pool.close(session_id))
|
||||
@@ -0,0 +1,348 @@
|
||||
"""
|
||||
Слоты браузера и окна сессий воркера (#81, #82).
|
||||
|
||||
Слот — один открытый Chrome. `BROWSER_SLOTS` ограничивает число окон воркера: и тех, где идёт задача,
|
||||
и удерживаемых между задачами. Окно сессии остаётся открытым после задачи, если в ней задан
|
||||
`browser_idle_timeout > 0`: следующая задача с той же cookie `session_id` попадает в это окно
|
||||
(маршрутизация — в api, #83) и продолжает работу без перезапуска Chrome. Время простоя считает
|
||||
Redis: TTL ключа `session:<id>`, каждая задача в окно ставит его заново. Ключ исчез — окно
|
||||
закрывается фоновой проверкой (`expire_idle`).
|
||||
|
||||
Задачи одной сессии выполняются строго по очереди: профиль Chrome нельзя открыть дважды.
|
||||
|
||||
Всё здесь работает в одном event loop (поток AsyncIO-middleware dramatiq), поэтому счётчик слотов
|
||||
и словарь окон не требуют блокировок между потоками.
|
||||
"""
|
||||
import asyncio
|
||||
import base64
|
||||
import logging
|
||||
import time
|
||||
from contextlib import asynccontextmanager, suppress
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from patchright.async_api import Page, ProxySettings
|
||||
|
||||
from api import actions, profile
|
||||
from api.geoip import resolve_geo
|
||||
from api.logs import safe_request_dump, safe_response_dump
|
||||
from api.schemas import TaskRequest
|
||||
from engine import AsyncStealthySession
|
||||
from tasks.store import RUNNING_TTL, Status, TaskStore
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
USER_DATA_ROOT = (Path(__file__).parent / "../../extra/user_data_dir").resolve()
|
||||
|
||||
|
||||
class Outcome:
|
||||
DONE = "done"
|
||||
# свободного слота нет — задачу нужно вернуть в очередь
|
||||
NO_SLOT = "no_slot"
|
||||
# задачу выполнять не нужно: отменена, уже выполнена или выполняется
|
||||
SKIPPED = "skipped"
|
||||
|
||||
|
||||
class NoWindow(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class Window:
|
||||
session_id: str
|
||||
engine: Any = None
|
||||
# прокси, с которым запущен Chrome: он задаётся при запуске, и задача с другим прокси перезапускает окно
|
||||
proxy: str | None = None
|
||||
# окно закрыть, как только закончится текущая задача (DELETE /v2/sessions/current)
|
||||
closing: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class _SessionLock:
|
||||
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
|
||||
users: int = 0
|
||||
|
||||
|
||||
def parse_proxy(proxy_url: str) -> ProxySettings:
|
||||
parts = urlsplit(proxy_url)
|
||||
settings: ProxySettings = {"server": f"{parts.scheme}://{parts.hostname}:{parts.port}"}
|
||||
if parts.username:
|
||||
settings["username"] = parts.username
|
||||
if parts.password:
|
||||
settings["password"] = parts.password
|
||||
return settings
|
||||
|
||||
|
||||
def session_dir(session_id: str) -> Path:
|
||||
user_data_dir = (USER_DATA_ROOT / session_id).resolve()
|
||||
if user_data_dir.parent != USER_DATA_ROOT:
|
||||
raise ValueError("invalid session path")
|
||||
return user_data_dir
|
||||
|
||||
|
||||
def make_page_action(payload: TaskRequest):
|
||||
"""page_action для AsyncStealthySession.fetch()/run(): выполняет очередь действий (вместе с
|
||||
прохождением антибота, api.actions.execute) и достаёт со страницы всё, что нужно для ответа.
|
||||
|
||||
Непройденный антибот не бросается наружу, а возвращается в challenge_error — чтобы в ответ
|
||||
всё равно попали cookies, html и скриншот страницы-заглушки."""
|
||||
|
||||
async def page_action(page: Page) -> dict:
|
||||
challenge_error = None
|
||||
try:
|
||||
actions_result = await actions.execute(
|
||||
page, payload.actions, payload.captcha,
|
||||
antibot_name=payload.captcha_type,
|
||||
success_locator=payload.captcha_success_locator,
|
||||
detect_locator=payload.captcha_detect_locator,
|
||||
)
|
||||
except actions.ChallengeFailed as e:
|
||||
challenge_error = str(e)
|
||||
actions_result = []
|
||||
|
||||
screenshot = None
|
||||
if payload.screenshot:
|
||||
screenshot = base64.b64encode(await page.screenshot(type="png")).decode()
|
||||
|
||||
return {
|
||||
"challenge_error": challenge_error,
|
||||
"actions": actions_result,
|
||||
"cookies": await page.context.cookies(),
|
||||
"html": await page.inner_html("html"),
|
||||
"screenshot": screenshot,
|
||||
"user_agent": await page.evaluate("() => navigator.userAgent"),
|
||||
}
|
||||
|
||||
return page_action
|
||||
|
||||
|
||||
def build_result(payload: TaskRequest, result: dict) -> dict:
|
||||
"""Ответ в формате SolveResponse из результата page_action."""
|
||||
failed = [a for a in result["actions"] if a.get("ok") is False]
|
||||
response = {
|
||||
"status": "ok",
|
||||
"user_agent": result["user_agent"],
|
||||
"cookies": result["cookies"],
|
||||
"actions": result["actions"],
|
||||
}
|
||||
if result["challenge_error"]:
|
||||
response.update(status="err", code="challenge_failed", error=result["challenge_error"])
|
||||
elif failed:
|
||||
response.update(status="err", code="action_failed", error=f"{failed[0]['type']}: {failed[0].get('result')}")
|
||||
if payload.html:
|
||||
response["html"] = result["html"]
|
||||
if payload.screenshot:
|
||||
response["screenshot"] = result["screenshot"]
|
||||
return response
|
||||
|
||||
|
||||
class WindowPool:
|
||||
def __init__(self, store: TaskStore, worker_id: str, slots: int):
|
||||
self.store = store
|
||||
self.worker_id = worker_id
|
||||
self.slots = slots
|
||||
self.used = 0
|
||||
self.windows: dict[str, Window] = {}
|
||||
self.running: set[str] = set()
|
||||
self._locks: dict[str, _SessionLock] = {}
|
||||
|
||||
def info(self) -> dict:
|
||||
return {
|
||||
"browser_slots": self.slots,
|
||||
"free_slots": self.slots - self.used,
|
||||
"windows": len(self.windows),
|
||||
"running": len(self.running),
|
||||
}
|
||||
|
||||
@asynccontextmanager
|
||||
async def _session_lock(self, session_id: str):
|
||||
entry = self._locks.setdefault(session_id, _SessionLock())
|
||||
entry.users += 1
|
||||
try:
|
||||
async with entry.lock:
|
||||
yield
|
||||
finally:
|
||||
entry.users -= 1
|
||||
if not entry.users:
|
||||
del self._locks[session_id]
|
||||
|
||||
async def report(self) -> None:
|
||||
"""heartbeat со слотами сразу после их изменения: по нему api решает, отвечать ли 503 на /v1/solve."""
|
||||
try:
|
||||
await self.store.heartbeat(self.worker_id, self.info())
|
||||
except Exception:
|
||||
log.warning("heartbeat failed", exc_info=True)
|
||||
|
||||
def _busy(self, session_id: str) -> bool:
|
||||
entry = self._locks.get(session_id)
|
||||
return bool(entry and entry.lock.locked())
|
||||
|
||||
async def handle(self, task_id: str, session_id: str, payload: dict, idle_timeout: int) -> str:
|
||||
"""Выполнить задачу. Результат — в Redis; возврат — что делать с сообщением очереди."""
|
||||
async with self._session_lock(session_id):
|
||||
window = self.windows.get(session_id)
|
||||
new_window = window is None
|
||||
if new_window:
|
||||
if self.used >= self.slots:
|
||||
task = await self.store.get(task_id)
|
||||
# отменённую или истёкшую задачу незачем возвращать в очередь
|
||||
return Outcome.NO_SLOT if task and task["status"] == Status.QUEUED else Outcome.SKIPPED
|
||||
# слот занимается до первого await, иначе его может перехватить другая задача
|
||||
self.used += 1
|
||||
window = self.windows[session_id] = Window(session_id)
|
||||
await self.report()
|
||||
|
||||
if not await self.store.claim(task_id, self.worker_id, running_here=task_id in self.running):
|
||||
if new_window:
|
||||
await self._close(window)
|
||||
return Outcome.SKIPPED
|
||||
|
||||
self.running.add(task_id)
|
||||
try:
|
||||
await self.store.set_session(session_id, self.worker_id, idle_timeout, RUNNING_TTL)
|
||||
result = await self._execute(window, TaskRequest.model_validate(payload))
|
||||
await self.store.finish(task_id, result)
|
||||
except Exception as e:
|
||||
log.exception("task %s failed", task_id)
|
||||
with suppress(Exception):
|
||||
await self.store.fail(task_id, str(e) or type(e).__name__)
|
||||
except asyncio.CancelledError:
|
||||
# dramatiq прервал задачу по time limit или при остановке воркера
|
||||
with suppress(Exception):
|
||||
await self.store.fail(task_id, "task interrupted")
|
||||
raise
|
||||
finally:
|
||||
self.running.discard(task_id)
|
||||
if idle_timeout and window.engine and not window.closing:
|
||||
await self.store.set_session(session_id, self.worker_id, idle_timeout, idle_timeout * 60)
|
||||
else:
|
||||
await self._close(window)
|
||||
with suppress(Exception):
|
||||
await self.store.delete_session(session_id)
|
||||
return Outcome.DONE
|
||||
|
||||
async def _execute(self, window: Window, payload: TaskRequest) -> dict:
|
||||
started = time.monotonic()
|
||||
session_id = window.session_id
|
||||
log.info("Incoming request [%s]: %s", session_id, payload.url or "<current page>")
|
||||
log.debug("Incoming payload [%s]:\n%s", session_id, safe_request_dump(payload))
|
||||
|
||||
page_action = make_page_action(payload)
|
||||
nav_kwargs = dict(
|
||||
timeout=payload.timeout * 1000 if payload.timeout else 60_000,
|
||||
disable_resources=payload.disable_resources,
|
||||
blocked_domains=set(payload.blocked_domains) if payload.blocked_domains else None,
|
||||
)
|
||||
try:
|
||||
if window.engine and window.proxy != payload.proxy:
|
||||
log.info("proxy changed, restarting browser [%s]", session_id)
|
||||
await self._close_engine(window)
|
||||
|
||||
if window.engine is None:
|
||||
if not payload.url:
|
||||
raise NoWindow("url is required: the session has no open browser window")
|
||||
window.engine = await self._open(session_id, payload, nav_kwargs)
|
||||
window.proxy = payload.proxy
|
||||
elif payload.cookies is not None:
|
||||
# переданные cookies заменяют cookies профиля, как и при запуске окна
|
||||
await window.engine.context.clear_cookies()
|
||||
if payload.cookies:
|
||||
await window.engine.context.add_cookies(payload.cookies)
|
||||
|
||||
if payload.url:
|
||||
raw = await window.engine.fetch(
|
||||
payload.url, page_action=page_action,
|
||||
wait=payload.wait, google_search=payload.google_search, **nav_kwargs,
|
||||
)
|
||||
else:
|
||||
raw = await window.engine.run(page_action, **nav_kwargs)
|
||||
except Exception as e:
|
||||
if not isinstance(e, NoWindow):
|
||||
log.exception("request failed [%s]", session_id)
|
||||
# окно могло остаться в непонятном состоянии: следующая задача откроет его заново из профиля
|
||||
await self._close_engine(window)
|
||||
result = {"status": "err", "code": "request_failed", "error": str(e) or type(e).__name__}
|
||||
else:
|
||||
result = build_result(payload, raw)
|
||||
|
||||
result["duration_ms"] = int((time.monotonic() - started) * 1000)
|
||||
log.info("Outgoing [%s]: %s %s %dms", session_id, result["status"], result.get("error", ""), result["duration_ms"])
|
||||
log.debug("Outgoing payload:\n%s", safe_response_dump(result))
|
||||
return result
|
||||
|
||||
async def _open(self, session_id: str, payload: TaskRequest, nav_kwargs: dict):
|
||||
user_data_dir = session_dir(session_id)
|
||||
user_data_dir.mkdir(parents=True, exist_ok=True)
|
||||
# отпечаток задаётся профилю один раз (см. api/profile.py)
|
||||
profile_kwargs = profile.browser_kwargs(profile.load_or_create(user_data_dir))
|
||||
|
||||
proxy = parse_proxy(payload.proxy) if payload.proxy else None
|
||||
# geo — по exit-IP прокси (см. api/geoip.py); без прокси оставляем настройки браузера
|
||||
geo = await resolve_geo(proxy) if proxy else None
|
||||
|
||||
engine = AsyncStealthySession(
|
||||
headless=False,
|
||||
proxy=proxy,
|
||||
user_data_dir=str(user_data_dir),
|
||||
real_chrome=True,
|
||||
locale=geo.get("locale") if geo else None,
|
||||
timezone_id=geo.get("timezone_id") if geo else None,
|
||||
cookies=payload.cookies,
|
||||
# переданные cookies заменяют cookies профиля, а не дополняют их
|
||||
clear_cookies=payload.cookies is not None,
|
||||
retries=max(1, payload.retries),
|
||||
retry_delay=payload.retry_delay,
|
||||
wait=payload.wait,
|
||||
google_search=payload.google_search,
|
||||
**nav_kwargs,
|
||||
**profile_kwargs,
|
||||
)
|
||||
await engine.start()
|
||||
return engine
|
||||
|
||||
async def _close_engine(self, window: Window) -> None:
|
||||
engine, window.engine = window.engine, None
|
||||
if engine is not None:
|
||||
try:
|
||||
await engine.close()
|
||||
except Exception:
|
||||
log.warning("failed to close browser [%s]", window.session_id, exc_info=True)
|
||||
|
||||
async def _close(self, window: Window) -> None:
|
||||
"""Закрыть окно и освободить его слот."""
|
||||
await self._close_engine(window)
|
||||
if self.windows.get(window.session_id) is window:
|
||||
del self.windows[window.session_id]
|
||||
self.used -= 1
|
||||
log.info("window closed [%s]", window.session_id)
|
||||
await self.report()
|
||||
|
||||
async def close(self, session_id: str) -> None:
|
||||
"""Закрыть окно сессии досрочно. Если в нём идёт задача — сразу после неё."""
|
||||
window = self.windows.get(session_id)
|
||||
if window is None:
|
||||
return
|
||||
window.closing = True
|
||||
if self._busy(session_id):
|
||||
return
|
||||
async with self._session_lock(session_id):
|
||||
if self.windows.get(session_id) is window:
|
||||
await self._close(window)
|
||||
|
||||
async def expire_idle(self) -> None:
|
||||
"""Закрыть окна, чей ключ сессии в Redis истёк (простой дольше browser_idle_timeout)."""
|
||||
for session_id, window in list(self.windows.items()):
|
||||
if self._busy(session_id) or await self.store.get_session(session_id):
|
||||
continue
|
||||
async with self._session_lock(session_id):
|
||||
if self.windows.get(session_id) is window and not await self.store.get_session(session_id):
|
||||
log.info("browser idle timeout [%s]", session_id)
|
||||
await self._close(window)
|
||||
|
||||
async def close_all(self) -> None:
|
||||
for session_id, window in list(self.windows.items()):
|
||||
await self._close(window)
|
||||
with suppress(Exception):
|
||||
await self.store.delete_session(session_id)
|
||||
Reference in new issue
Block a user