""" Сервис 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 logging import os import time import uuid from typing import Annotated from fastapi import Cookie, FastAPI, HTTPException, Response from fastapi.responses import JSONResponse 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.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__) RETRY_AFTER = os.getenv("NO_SLOTS_RETRY_AFTER", "5") # сколько /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 "") 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(): 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, ): """Синхронный контракт 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() 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) result.setdefault("actions", []) result["duration_ms"] = int((time.monotonic() - started) * 1000) return result