diff --git a/Dockerfile b/Dockerfile index 7d37b54..5f1c049 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,3 +1,4 @@ +# Сервис worker (#78): выполняет задачи из очереди в Google Chrome под Xvfb. FROM python:3.12-slim-bookworm ENV PYTHONDONTWRITEBYTECODE=1 \ @@ -30,9 +31,9 @@ WORKDIR /app/src # профили браузера (extra/user_data_dir) и geoip-кэш — в volume VOLUME ["/app/extra"] -EXPOSE 8000 # tini как PID 1: пробрасывает сигналы и подбирает осиротевшие процессы Chrome/Xvfb (иначе — зомби) ENTRYPOINT ["/usr/bin/tini", "-g", "--"] -CMD ["xvfb-run", "-a", "--server-args=-screen 0 1920x1080x24", \ - "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"] +# один процесс: открытые окна браузера живут в его памяти; потоков — с запасом над слотами, чтобы +# задачи в уже открытые окна не ждали, пока потоки заняты задачами без свободного слота +CMD ["sh", "-c", "exec xvfb-run -a --server-args='-screen 0 1920x1080x24' dramatiq worker.actors --processes 1 --threads ${WORKER_THREADS:-$((${BROWSER_SLOTS:-4} * 2))}"] diff --git a/Dockerfile.api b/Dockerfile.api new file mode 100644 index 0000000..749a2cf --- /dev/null +++ b/Dockerfile.api @@ -0,0 +1,23 @@ +# Сервис api (#78): принимает задачи и отдаёт результаты. Без Chrome — лёгкий образ. +FROM python:3.12-slim-bookworm + +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 + +WORKDIR /app + +COPY requirements-api.txt . +RUN pip install -r requirements-api.txt \ + && useradd --create-home --uid 1000 app + +COPY --chown=app:app src/main.py ./src/ +COPY --chown=app:app src/tasks/ ./src/tasks/ +COPY --chown=app:app src/api/__init__.py src/api/schemas.py ./src/api/ + +USER app +WORKDIR /app/src +EXPOSE 8000 + +CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"] diff --git a/Makefile b/Makefile index a730bb9..73384d4 100644 --- a/Makefile +++ b/Makefile @@ -5,10 +5,16 @@ WHEELS_DIR := wheels REQUIREMENTS := requirements.txt ver ?= latest -MAKECMDGOALS_TARGETS := install freeze wheels sync sync-offline clean help build release sync-dev test test-browser +MAKECMDGOALS_TARGETS := install freeze wheels sync sync-offline clean help build release sync-dev test test-browser up down logs smoke PKG := $(filter-out $(MAKECMDGOALS_TARGETS),$(MAKECMDGOALS)) -.PHONY: install freeze wheels sync sync-offline clean help venv build release sync-dev test test-browser +.PHONY: install freeze wheels sync sync-offline clean help venv build release sync-dev test test-browser up down logs smoke + +# образы сервисов: api — без Chrome, worker — Chrome + Xvfb; rabbitmq и redis — сторонние, кладутся в релиз +API_IMAGE := robotex-api +WORKER_IMAGE := robotex-worker +DEPS_IMAGES := rabbitmq:4.1-alpine redis:8.2-alpine +API_URL ?= http://localhost:8000 help: @echo "make install [ ...] - установить пакет(ы) в $(VENV), обновить $(REQUIREMENTS), собрать wheel в $(WHEELS_DIR)/" @@ -17,8 +23,12 @@ help: @echo "make freeze - перезаписать $(REQUIREMENTS) текущим состоянием окружения" @echo "make wheels - собрать wheel для всех пакетов из $(REQUIREMENTS)" @echo "make clean - удалить $(WHEELS_DIR)/" - @echo "make build [ver=latest] - собрать docker-образ robotex:" - @echo "make release ver= - собрать образ и пакет для деплоя в build// (compose, .env, образ .tar.gz)" + @echo "make build [ver=latest] - собрать образы $(API_IMAGE): и $(WORKER_IMAGE):" + @echo "make release ver= - собрать образы и пакет для деплоя в build// (compose, .env, образы .tar.gz)" + @echo "make up - собрать и запустить локально: api, worker, rabbitmq, redis" + @echo "make down - остановить локальный стенд" + @echo "make logs - логи api и worker" + @echo "make smoke [API_URL=...] - проверить запущенный стенд: /health, /v1/solve, /v2/tasks, удержание окна" @echo "make sync-dev - установить зависимости для тестов (requirements-dev.txt)" @echo "make test - юнит-тесты без браузера" @echo "make test-browser [ver=latest] - собрать образ и прогнать сквозные тесты с Chrome в Docker (нужна сеть)" @@ -53,8 +63,20 @@ clean: rm -rf $(WHEELS_DIR) build: - @echo "build an application" - docker build -t robotex:$(ver) . + docker build -f Dockerfile.api -t $(API_IMAGE):$(ver) . + docker build -t $(WORKER_IMAGE):$(ver) . + +up: build + docker compose up -d + +down: + docker compose down + +logs: + docker compose logs -f api worker + +smoke: + $(PYTHON) tests/smoke.py $(API_URL) sync-dev: venv $(PIP) install -r requirements-dev.txt @@ -66,7 +88,7 @@ test-browser: build docker run --rm --shm-size=2g \ -v $(CURDIR)/tests:/app/tests:ro -v $(CURDIR)/pytest.ini:/app/pytest.ini:ro \ -v $(CURDIR)/requirements-dev.txt:/app/requirements-dev.txt:ro \ - -w /app robotex:$(ver) sh -c 'pip install -q --user -r requirements-dev.txt && \ + -w /app $(WORKER_IMAGE):$(ver) sh -c 'pip install -q --user -r requirements-dev.txt && \ xvfb-run -a --server-args="-screen 0 1920x1080x24" python -m pytest -m browser -p no:cacheprovider -s' RELEASE_DIR := build/$(ver) @@ -80,9 +102,11 @@ endif release: build @echo "make release $(ver) -> $(RELEASE_DIR)/" @mkdir -p $(RELEASE_DIR) - sed 's|image: robotex:.*|image: robotex:$(ver)|' docker-compose.prod.yml > $(RELEASE_DIR)/docker-compose.yml + sed -e 's|image: $(API_IMAGE):.*|image: $(API_IMAGE):$(ver)|' -e 's|image: $(WORKER_IMAGE):.*|image: $(WORKER_IMAGE):$(ver)|' \ + docker-compose.prod.yml > $(RELEASE_DIR)/docker-compose.yml cp dot_env $(RELEASE_DIR)/.env - docker save robotex:$(ver) | gzip > $(RELEASE_DIR)/robotex-$(ver).tar.gz + for img in $(DEPS_IMAGES); do docker image inspect $$img >/dev/null 2>&1 || docker pull $$img; done + docker save $(API_IMAGE):$(ver) $(WORKER_IMAGE):$(ver) $(DEPS_IMAGES) | gzip > $(RELEASE_DIR)/robotex-$(ver).tar.gz @printf '%s\n' \ '# Robotex $(ver)' \ '' \ @@ -91,7 +115,7 @@ release: build 'docker compose up -d' \ '```' \ '' \ - 'Настройки (порт, CPU, память) — в .env.' > $(RELEASE_DIR)/README.md + 'Настройки (порт, слоты, CPU и память воркера, пароль RabbitMQ) — в .env.' > $(RELEASE_DIR)/README.md @echo "Готово: $(RELEASE_DIR)/" # позволяет передавать имя пакета как аргумент: make install redis diff --git a/README.md b/README.md index 2bc762d..4f09ad6 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ behind Cloudflare-style JS challenges. ## Status -Early MVP (v0), under active development. Single browser engine, no queue, no billing yet. Tasks and project wiki — +Early MVP (v0), under active development. Single browser engine, no billing yet. Tasks and project wiki — [Redmine](https://redmine.sokolab.xyz/projects/robotex). Not implemented yet, though planned for v1: @@ -18,22 +18,25 @@ Not implemented yet, though planned for v1: ## Architecture ``` -Client → Caddy (TLS, on the host) → FastAPI POST /v1/solve → browser slot (semaphore, 503 if none) - │ - ▼ - Chrome (Patchright), profile per session_id - pass anti-bot → run actions → collect result - │ - ▼ - release slot → response +Client → Caddy (TLS) → api (FastAPI) ──task──▶ RabbitMQ ──▶ worker (dramatiq, Chrome + Xvfb) + ▲ │ POST /v2/tasks common queue browser slots, open windows + │ │ POST /v1/solve or the worker's pass anti-bot → run actions + │ ▼ own queue │ + └─── GET /v2/tasks/{id} ◀── Redis: task status and result (TTL) ◀─┘ ``` - **Engine**: real Google Chrome (`channel="chrome"`) driven via [Patchright](https://github.com/Kaliiiiiiiiii-Vinyzu/patchright-python) (undetected Playwright fork), headful under Xvfb. Browser session wrapper ported from scrapling. Firefox is not supported yet. -- **API**: FastAPI, single versioned endpoint (`POST /v1/solve`) plus `GET /health`. -- **Concurrency**: bounded by an `asyncio.Semaphore` of `BROWSER_SLOTS` (env, default 10) — no external job queue in - v1. When all slots are busy the request gets `503` with `Retry-After` immediately, it is not queued. +- **Two services**: `api` (FastAPI, no browser) registers tasks and serves results; `worker` (dramatiq) runs them in + Chrome. Tasks go through RabbitMQ, statuses and results live in Redis with a TTL. +- **API**: `POST /v2/tasks` + `GET /v2/tasks/{id}` (asynchronous, poll by link), `DELETE /v2/sessions/current`; + `POST /v1/solve` keeps the synchronous contract on top of the same queue; `GET /health`. +- **Concurrency**: each worker has `BROWSER_SLOTS` Chrome windows. Tasks without a free slot wait in the queue; + `/v1/solve` answers `503` with `Retry-After` when no slot is free. +- **Window holding**: with `browser_idle_timeout` (minutes, 0–120) the session's Chrome stays open after the task; + the next task with the same `session_id` is routed to that worker's own queue and runs in the open window, even + without a `url` (on the current page). Idle time is a Redis key TTL; an open window occupies a slot. - **Anti-bot**: detect → solve → verify loop for Cloudflare Turnstile and reCAPTCHA; captchas inside pages (Turnstile, reCAPTCHA, hCaptcha) are solved via the client's 2captcha key. - **Proxies**: not managed by the service — the client may supply their own proxy per request. Browser locale and @@ -44,19 +47,27 @@ Client → Caddy (TLS, on the host) → FastAPI POST /v1/solve → browser slot mismatches `userAgentData` and Cloudflare blocks it. Clients take `user_agent` from the response. - **Logging**: stdout, level from `LOG_LEVEL`; client secrets (proxy credentials, 2captcha key, `fill` values, cookie values) are masked. -- **Deployment**: single Docker container; TLS is terminated by Caddy on the host — see +- **Deployment**: Docker Compose — `api`, `worker`, `rabbitmq`, `redis`; TLS is terminated by Caddy on the host — see [docs/Deployment.md](docs/Deployment.md). ### Project layout ``` src/ - main.py # FastAPI app: /v1/solve, /health, slots, sessions, response + main.py # api service: /v2/tasks, /v1/solve, /health, sessions + tasks/ # shared by api and worker, no browser + store.py # Redis: tasks, results, open windows, worker heartbeats + broker.py # dramatiq broker (RabbitMQ), queues, messages + sessions.py # session_id cookie + worker/ + actors.py # dramatiq actors, worker lifecycle, heartbeat + windows.py # browser slots, windows held between tasks, task execution api/ - profile.py # per-profile user agent and generated fingerprint - schemas.py # SolveRequest / ActionRequest / SolveResponse + validation + profile.py # per-profile generated fingerprint + schemas.py # SolveRequest / TaskRequest / ActionRequest / responses + validation actions.py # action registry and execution geoip.py # proxy → locale / timezone + logs.py # request/response dumps without client secrets engine/ # Chrome session wrapper (ported from scrapling) stealthy.py # AsyncStealthySession session.py # browser / context launch options @@ -71,10 +82,13 @@ docs/ # API, deployment, anti-bot notes tests/ unit/ # pytest without a browser: make test browser/ # end-to-end with real Chrome in Docker: make test-browser + harness.py # api + worker in one process (fakeredis, no RabbitMQ) for tests + smoke.py # check a running stack: make smoke loadtest.py # concurrent load test + docker stats test_main.http # manual requests (JetBrains HTTP client) wheels/ # wheels for offline install on Ubuntu Noble -Dockerfile +Dockerfile # worker: Chrome + Xvfb +Dockerfile.api # api: no browser docker-compose.yml # local / load testing docker-compose.prod.yml # template for releases (make release) Makefile @@ -82,6 +96,8 @@ Makefile ## API +Synchronous: + ``` POST /v1/solve X-API-Key: @@ -111,15 +127,27 @@ Response: ``` On failure `status` is `"err"` with a machine-readable `code` (`no_slots`, `request_failed`, `challenge_failed`, -`action_failed`) and a human-readable `error`. Full reference — fields, action types, errors, sessions: +`action_failed`) and a human-readable `error`. + +Asynchronous, keeping the browser window open between tasks: + +``` +POST /v2/tasks {"url": "https://example.com/login", "actions": [...], "browser_idle_timeout": 10} +→ 202 {"task_id": "...", "status": "queued", "url": "/v2/tasks/..."} + Set-Cookie: session_id +GET /v2/tasks/{id} → {"status": "queued|running|done|failed", "result": {... same as /v1/solve ...}} +POST /v2/tasks {"actions": [...], "browser_idle_timeout": 10} # same session: runs in the open window +DELETE /v2/sessions/current # close the window, free the slot +``` + +Full reference — fields, action types, errors, sessions, v2: [docs/API.md](docs/API.md). How to use the API from your client (same UA and proxy as the browser, sessions, retries) — [docs/Usage.md](docs/Usage.md). ## Roadmap -Deliberately out of scope for v1: +Not done yet: -- Job queue (Redis/RabbitMQ) + multiple worker VPS for horizontal scaling +- Workers on several hosts (needs shared or per-host routing of browser profiles) - Real billing/usage plans beyond a flat API-key check - Broader action vocabulary - Second engine (Firefox, e.g. Camoufox) for targets where the Chrome fingerprint doesn't fit @@ -130,33 +158,40 @@ Deliberately out of scope for v1: python -m venv .venv source .venv/bin/activate make sync-offline -cd src && uvicorn main:app --reload +docker compose up -d rabbitmq redis # queue and storage +cd src && RABBITMQ_URL=... REDIS_URL=... uvicorn main:app --reload +cd src && dramatiq worker.actors --processes 1 --threads 4 ``` -Chrome runs headful, so a display is required (locally — your desktop, in Docker — Xvfb). +The compose network is not published, so for running the services on the host publish `5672` and `6379` yourself, +or just use `make up`. Chrome runs headful, so a display is required (locally — your desktop, in Docker — Xvfb). Tests: ```bash make sync-dev # pytest into .venv -make test # unit tests, no browser (~5 s) -make test-browser # build the image and run end-to-end tests with real Chrome (needs network) +make test # unit tests: api + worker + dramatiq, no browser, no RabbitMQ/Redis (~10 s) +make test-browser # build the images and run end-to-end tests with real Chrome (needs network) +make smoke # check the running stack from make up ``` ## Docker ```bash -docker compose up -d --build +make up # build robotex-api and robotex-worker, start api, worker, rabbitmq, redis +make smoke # health, /v1/solve, /v2/tasks, window holding +make logs +make down ``` The API listens on `http://localhost:8000`. Browser profiles (`extra/user_data_dir`) and the GeoIP cache live in the -`extra` named volume, so sessions survive container restarts. +worker's `extra` named volume, so sessions survive container restarts. Releases: ```bash -make build ver=v0.0.2 # docker image robotex:v0.0.2 -make release ver=v0.0.2 # image + deploy bundle in build/v0.0.2/ (compose, .env, image .tar.gz) +make build ver=v0.1.0 # docker images robotex-api:v0.1.0 and robotex-worker:v0.1.0 +make release ver=v0.1.0 # images + deploy bundle in build/v0.1.0/ (compose, .env, images .tar.gz) ``` Details (Xvfb, tini, `/dev/shm`, env settings, load measurements) — [docs/Deployment.md](docs/Deployment.md). diff --git a/docker-compose.prod.yml b/docker-compose.prod.yml index 7b44c80..1cae2ed 100644 --- a/docker-compose.prod.yml +++ b/docker-compose.prod.yml @@ -1,22 +1,71 @@ +x-backend-env: &backend-env + RABBITMQ_URL: amqp://robotex:${RABBITMQ_PASSWORD:-robotex}@rabbitmq:5672/ + REDIS_URL: redis://redis:6379/0 + services: api: - image: robotex:v0.0.1 + image: robotex-api:latest restart: unless-stopped ports: - "${API_PORT:-8000}:8000" env_file: - path: .env required: false + environment: *backend-env + depends_on: + rabbitmq: {condition: service_healthy} + redis: {condition: service_healthy} + + worker: + image: robotex-worker:latest + restart: unless-stopped + env_file: + - path: .env + required: false + environment: *backend-env volumes: - api-data:/app/extra # Chrome активно использует /dev/shm, дефолтных 64MB не хватает shm_size: "2gb" # Потолок ресурсов: при переборе OOM-killer прибьёт процесс Chrome внутри контейнера, # а не соседние сервисы на хосте. Подбирать под сервер вместе с BROWSER_SLOTS. - cpus: ${API_CPUS:-4} - mem_limit: ${API_MEM_LIMIT:-8g} - memswap_limit: ${API_MEM_LIMIT:-8g} + cpus: ${WORKER_CPUS:-4} + mem_limit: ${WORKER_MEM_LIMIT:-8g} + memswap_limit: ${WORKER_MEM_LIMIT:-8g} + # при остановке воркер дожидается текущих задач и закрывает окна + stop_grace_period: 60s + depends_on: + rabbitmq: {condition: service_healthy} + redis: {condition: service_healthy} + rabbitmq: + image: rabbitmq:4.1-alpine + restart: unless-stopped + environment: + RABBITMQ_DEFAULT_USER: robotex + RABBITMQ_DEFAULT_PASS: ${RABBITMQ_PASSWORD:-robotex} + # задача выполняется до 3 ч, а по умолчанию RabbitMQ рвёт канал, если сообщение не подтверждено 30 мин + RABBITMQ_SERVER_ADDITIONAL_ERL_ARGS: -rabbit consumer_timeout 14400000 + volumes: + - rabbitmq:/var/lib/rabbitmq + healthcheck: + test: ["CMD", "rabbitmq-diagnostics", "-q", "ping"] + interval: 10s + timeout: 10s + retries: 10 + + redis: + image: redis:8.2-alpine + restart: unless-stopped + volumes: + - redis:/data + healthcheck: + test: ["CMD", "redis-cli", "ping"] + interval: 5s + timeout: 5s + retries: 10 volumes: api-data: + rabbitmq: + redis: diff --git a/docker-compose.yml b/docker-compose.yml index e259f0b..42049b1 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,25 +1,78 @@ +# Локальный запуск и тесты: docker compose up -d --build +# Шаблон для релизов — docker-compose.prod.yml (make release). +x-backend-env: &backend-env + RABBITMQ_URL: amqp://robotex:${RABBITMQ_PASSWORD:-robotex}@rabbitmq:5672/ + REDIS_URL: redis://redis:6379/0 + services: api: - build: . - image: robotex:latest + build: + context: . + dockerfile: Dockerfile.api + image: robotex-api:latest restart: unless-stopped ports: - "${API_PORT:-8000}:8000" env_file: - path: .env required: false + environment: *backend-env + depends_on: + rabbitmq: {condition: service_healthy} + redis: {condition: service_healthy} + + worker: + build: . + image: robotex-worker:latest + restart: unless-stopped + env_file: + - path: .env + required: false + environment: *backend-env volumes: - extra:/app/extra # Chrome активно использует /dev/shm, дефолтных 64MB не хватает shm_size: "2gb" # Потолок ресурсов: при переборе OOM-killer прибьёт процесс Chrome внутри контейнера, # а не соседние сервисы на хосте. Подбирать под сервер вместе с BROWSER_SLOTS. - cpus: ${API_CPUS:-4} - mem_limit: ${API_MEM_LIMIT:-8g} - memswap_limit: ${API_MEM_LIMIT:-8g} + cpus: ${WORKER_CPUS:-4} + mem_limit: ${WORKER_MEM_LIMIT:-8g} + memswap_limit: ${WORKER_MEM_LIMIT:-8g} + # при остановке воркер дожидается текущих задач и закрывает окна + stop_grace_period: 60s + depends_on: + rabbitmq: {condition: service_healthy} + redis: {condition: service_healthy} + + rabbitmq: + image: rabbitmq:4.1-alpine + restart: unless-stopped + environment: + RABBITMQ_DEFAULT_USER: robotex + RABBITMQ_DEFAULT_PASS: ${RABBITMQ_PASSWORD:-robotex} + # задача выполняется до 3 ч, а по умолчанию RabbitMQ рвёт канал, если сообщение не подтверждено 30 мин + RABBITMQ_SERVER_ADDITIONAL_ERL_ARGS: -rabbit consumer_timeout 14400000 + volumes: + - rabbitmq:/var/lib/rabbitmq + healthcheck: + test: ["CMD", "rabbitmq-diagnostics", "-q", "ping"] + interval: 10s + timeout: 10s + retries: 10 + + redis: + image: redis:8.2-alpine + restart: unless-stopped + volumes: + - redis:/data + healthcheck: + test: ["CMD", "redis-cli", "ping"] + interval: 5s + timeout: 5s + retries: 10 # HTTP-прокси для локальных тестов (exit-IP = IP хоста): docker compose --profile loadtest up -d - # из контейнера api доступен как http://proxy:8080 + # из контейнера worker доступен как http://proxy:8080 proxy: image: gogost/gost:3.1.0 profiles: ["loadtest"] @@ -28,3 +81,5 @@ services: volumes: extra: + rabbitmq: + redis: diff --git a/docs/API.md b/docs/API.md index d0d3da4..1c9be4f 100644 --- a/docs/API.md +++ b/docs/API.md @@ -2,11 +2,17 @@ Как встраивать API в свой клиент (UA, прокси, сессии, повторы) — [Usage.md](Usage.md). -Основной эндпоинт: `POST /v1/solve`. Открывает страницу в Chrome (через прокси клиента, если он задан), при необходимости проходит -антибот (Cloudflare Turnstile, reCAPTCHA), выполняет список действий и возвращает cookies, результаты действий и, -по запросу, HTML и скриншот. +Задача открывает страницу в Chrome (через прокси клиента, если он задан), при необходимости проходит антибот +(Cloudflare Turnstile, reCAPTCHA), выполняет список действий и возвращает cookies, результаты действий и, по запросу, +HTML и скриншот. Задачи выполняет сервис worker из очереди; API работает в двух режимах: -`GET /health` (и `GET /`) — `{"status": "ok", "browser_slots": N, "free_slots": M}`. +* `POST /v1/solve` — синхронно: ответ приходит, когда задача выполнена; +* `POST /v2/tasks` + `GET /v2/tasks/{id}` — асинхронно, с возможностью держать окно браузера открытым между + задачами (см. [API v2](#api-v2)). + +`GET /health` (и `GET /`) — `{"status": "ok", "browser_slots": N, "free_slots": M, "workers": [...]}`: слоты всех +живых воркеров; `workers` — по каждому `worker_id`, `browser_slots`, `free_slots`, `windows` (открытые окна), +`running` (выполняемые задачи). > `X-API-Key` пока не проверяется ([#4](https://redmine.sokolab.xyz/issues/4)). @@ -104,7 +110,7 @@ Cookie: session_id= # необязательно, см. «Сесси | `code` | HTTP | Когда | Что ещё есть в ответе | |---|---|---|---| -| `no_slots` | `503` + `Retry-After` | все браузеры заняты (`BROWSER_SLOTS`), запрос не ставится в очередь | — | +| `no_slots` | `503` + `Retry-After` | только `/v1/solve`: все слоты браузеров заняты, или воркер не взял задачу за `V1_QUEUE_WAIT` секунд (задача отменяется) | — | | `request_failed` | `200` | ошибка до действий: навигация, прокси, неизвестный `captcha_type` | `actions: []` | | `challenge_failed` | `200` | антибот не пройден за отведённое время, действия не выполнялись | cookies, html, screenshot | | `action_failed` | `200` | упало действие, выполнение остановлено на нём | cookies, html, screenshot, `actions` до упавшего включительно | @@ -139,9 +145,72 @@ canvas. Ограничение: canvas, WebGL и `hardwareConcurrency` у все User-Agent и размер окна запросом не задаются — это родные значения браузера (почему — [Usage.md](Usage.md#что-задаёт-сервис-а-не-клиент)). Поле `user_agent` в запросе даёт `422`. -Один профиль нельзя использовать в двух запросах одновременно: Chrome заблокирует каталог профиля. +Задачи одной сессии выполняются по очереди: профиль Chrome нельзя открыть дважды. Вторая задача ждёт, пока +закончится первая. + +## API v2 + +Асинхронный режим: задача ставится в очередь, клиент опрашивает её по ссылке. Тело задачи — как у `/v1/solve`, плюс: + +| Поле | Тип | По умолчанию | Описание | +|---|---|---|---| +| `browser_idle_timeout` | int, мин | `0` | держать окно браузера открытым после задачи столько минут простоя, `0`–`120`. `0` — закрыть сразу, как в v1 | +| `url` | string | `null` | необязателен, если окно сессии открыто: тогда действия выполняются на текущей странице, без навигации | + +### POST /v2/tasks + +```http +POST /v2/tasks +Content-Type: application/json +Cookie: session_id= # необязательно + +{"url": "https://freelancehunt.com/projects", "captcha_type": "cloudflare", "actions": [], "browser_idle_timeout": 10} +``` + +Ответ `202` и cookie `session_id`: + +```json +{"task_id": "0b9c…", "status": "queued", "url": "/v2/tasks/0b9c…"} +``` + +Без `url` и без открытого окна сессии — `422`. + +### GET /v2/tasks/{id} + +Только с cookie `session_id` той сессии, в которой задача создана; иначе, а также для несуществующей или истёкшей +задачи — `404`. + +```json +{"task_id": "0b9c…", "status": "done", "result": {"status": "ok", "user_agent": "…", "cookies": […], "actions": […], "duration_ms": 7120}} +``` + +| `status` | Значение | +|---|---| +| `queued` | ждёт свободного слота браузера | +| `running` | выполняется | +| `done` | готово, ответ в `result` — в том же формате, что ответ `/v1/solve` (со своим `status` `ok`/`err`) | +| `failed` | задачу не удалось выполнить (например, воркер упал дважды), причина в `error` | + +Результат хранится `RESULT_TTL_HOURS` часов (по умолчанию 24), затем `404`. Опрашивать раз в 0,5–2 секунды. + +### Окно браузера между задачами + +С `browser_idle_timeout > 0` Chrome сессии не закрывается после задачи. Следующая задача с той же cookie +`session_id` выполняется в этом окне: без запуска Chrome и обычно без повторного антибота, а без `url` — прямо на +текущей странице (например, продолжить заполнение формы). Каждая задача в окно заново отсчитывает простой — по +своему `browser_idle_timeout`; задача с `0` закрывает окно после себя. + +* Открытое окно занимает слот браузера, пока не закроется. Не держите окна дольше, чем нужно. +* Задача с другим `proxy`, чем у открытого окна, перезапускает Chrome (прокси задаётся при запуске). +* `cookies` в задаче заменяют cookies открытого окна, как и при запуске. +* Если задача в окне упала с `request_failed`, окно закрывается; следующая задача откроет его заново из профиля. + +### DELETE /v2/sessions/current + +Закрыть окно сессии (по cookie `session_id`) досрочно и освободить слот. Если в окне идёт задача — сразу после неё. +Ответ `{"closed": true}`, или `{"closed": false}`, если окна нет. ## Производительность -Каждый запрос запускает отдельный Chrome: около 0,6 ГБ памяти и до 3–4 ядер на пике при прохождении Cloudflare. +Каждое окно — отдельный Chrome: около 0,6 ГБ памяти и до 3–4 ядер на пике при прохождении Cloudflare. Замеры и рекомендации по лимитам — в [Deployment.md](Deployment.md#нагрузка). diff --git a/docs/Deployment.md b/docs/Deployment.md index 130fb4c..97f6dd1 100644 --- a/docs/Deployment.md +++ b/docs/Deployment.md @@ -3,20 +3,35 @@ ## Быстрый старт ```bash -docker compose up -d --build -docker compose logs -f api +make up # собрать образы и запустить api, worker, rabbitmq, redis +make smoke # проверить: /health, /v1/solve, /v2/tasks, удержание окна +make logs # логи api и worker +make down ``` -API доступно на `http://localhost:8000` (`POST /v1/solve`). +API доступно на `http://localhost:8000`. -Без compose: +## Сервисы -```bash -docker build -t robotex-api . -docker run -d --name robotex -p 8000:8000 --shm-size=2g -v robotex-extra:/app/extra robotex-api +``` +клиент ─▶ api ──▶ RabbitMQ ──▶ worker (Chrome + Xvfb) ──▶ Redis ◀── api ``` -## Что внутри образа +| Сервис | Образ | Что делает | +|---|---|---| +| `api` | `robotex-api` (`Dockerfile.api`, без Chrome) | принимает задачи, ставит в очередь, отдаёт результаты из Redis; `/v1/solve` ждёт результат сам | +| `worker` | `robotex-worker` (`Dockerfile`) | `dramatiq worker.actors`: выполняет задачи в Chrome, держит окна сессий (`browser_idle_timeout`) | +| `rabbitmq` | `rabbitmq:4.1-alpine` | очереди: общая `robotex` и персональная `robotex.worker.` у каждого воркера | +| `redis` | `redis:8.2-alpine` | задачи и результаты (TTL `RESULT_TTL_HOURS`), открытые окна (TTL = простой), heartbeat воркеров | + +Наружу открыт только порт `api`. Прокси и ключ 2captcha из запроса передаются только в сообщении очереди и в Redis +не сохраняются. + +Воркеров можно запустить несколько (`docker compose up -d --scale worker=2`): каждый берёт задачи из общей очереди, +а задачи сессии с открытым окном api направляет в персональную очередь воркера, где это окно. Профили сессий лежат +в volume воркера — воркеры на одном хосте делят его. + +## Что внутри образа worker | Компонент | Зачем | |---|---| @@ -27,11 +42,12 @@ docker run -d --name robotex -p 8000:8000 --shm-size=2g -v robotex-extra:/app/ex | `tini` (ENTRYPOINT, PID 1) | пробрасывает сигналы и подбирает осиротевшие процессы Chrome/Xvfb — без него копятся зомби-процессы | | пользователь `app` (uid 1000) | Chrome не запускается от root без `--no-sandbox` | -Рабочая директория — `/app/src`, uvicorn запускается как `uvicorn main:app` (импорты в проекте вида `from api import ...`). +Рабочая директория — `/app/src`. Воркер — один процесс dramatiq (`--processes 1`: открытые окна живут в его памяти), +потоков `WORKER_THREADS` (по умолчанию `2 × BROWSER_SLOTS`). ## Данные и volume -`/app/extra` вынесен в volume `extra`: +`/app/extra` воркера вынесен в volume (`extra` локально, `api-data` в релизе): - `extra/user_data_dir/` — профили браузера (cookies, local storage, fingerprint) по `session_id`; - `extra/geoip_cache.json` — кэш геолокации прокси. @@ -39,15 +55,18 @@ docker run -d --name robotex -p 8000:8000 --shm-size=2g -v robotex-extra:/app/ex Профили со временем накапливаются — чистить вручную: ```bash -docker compose exec api sh -c 'rm -rf /app/extra/user_data_dir/*' +docker compose exec worker sh -c 'rm -rf /app/extra/user_data_dir/*' ``` ## Настройки compose -- `shm_size: 2gb` — Chrome активно использует `/dev/shm`, дефолтных 64 МБ не хватает (вкладки падают с «Aw, Snap»). -- `API_PORT` — порт на хосте (по умолчанию 8000): `API_PORT=8001 docker compose up -d`. -- `API_CPUS` / `API_MEM_LIMIT` — потолок ресурсов контейнера (по умолчанию 4 CPU / 8 ГБ, swap выключен). +- `shm_size: 2gb` у worker — Chrome активно использует `/dev/shm`, дефолтных 64 МБ не хватает (вкладки падают с «Aw, Snap»). +- `API_PORT` — порт api на хосте (по умолчанию 8000): `API_PORT=8001 docker compose up -d`. +- `WORKER_CPUS` / `WORKER_MEM_LIMIT` — потолок ресурсов воркера (по умолчанию 4 CPU / 8 ГБ, swap выключен). При переборе памяти OOM-killer убивает процесс Chrome внутри контейнера, хост не страдает. +- `RABBITMQ_PASSWORD` — пароль пользователя `robotex` в RabbitMQ (по умолчанию `robotex`; порт наружу не открыт). + RabbitMQ запускается с `consumer_timeout` 4 ч: задача может выполняться дольше 30 минут по умолчанию. +- `stop_grace_period: 60s` у worker — при остановке он доделывает текущие задачи и закрывает окна. - `restart: unless-stopped`. - `.env` подключается, если существует (шаблон — `env.example`). @@ -55,8 +74,14 @@ docker compose exec api sh -c 'rm -rf /app/extra/user_data_dir/*' | Переменная | По умолчанию | Описание | |---|---|---| -| `BROWSER_SLOTS` | `10` | сколько Chrome работает параллельно; сверх — `503` | -| `NO_SLOTS_RETRY_AFTER` | `5` | значение `Retry-After` (с) в ответе `503` | +| `BROWSER_SLOTS` | `4` | worker: сколько окон Chrome открыто одновременно — и выполняющих задачу, и удерживаемых. Задача без свободного слота ждёт в очереди | +| `WORKER_THREADS` | `2 × BROWSER_SLOTS` | worker: потоков dramatiq | +| `WORKER_ID` | hostname | worker: имя воркера и его персональной очереди | +| `RESULT_TTL_HOURS` | `24` | сколько хранится результат задачи | +| `V1_QUEUE_WAIT` | `30` | api: сколько `/v1/solve` ждёт, пока воркер возьмёт задачу; дальше `503` | +| `NO_SLOTS_RETRY_AFTER` | `5` | api: значение `Retry-After` (с) в ответе `503` | +| `RABBITMQ_URL` | `amqp://guest:guest@localhost:5672/` | api, worker; в compose задаётся сам | +| `REDIS_URL` | `redis://localhost:6379/0` | api, worker; в compose задаётся сам | | `LOG_LEVEL` | `INFO` | `DEBUG` — полный дамп запроса/ответа (секреты клиента маскируются) | ## Нагрузка @@ -80,9 +105,12 @@ docker compose exec api sh -c 'rm -rf /app/extra/user_data_dir/*' время растёт в десятки раз, а Cloudflare перестаёт пропускать (на 4 CPU и 10 параллельных — 6 из 10 застряли на челлендже). - Память — около **0,5–0,6 ГБ** на браузер. -- Правило подбора: `BROWSER_SLOTS ≈ число ядер`, `API_MEM_LIMIT ≈ 0,6 ГБ × BROWSER_SLOTS + 1 ГБ`. - Пример: сервер 4 ядра / 8 ГБ → 3–4 слота. `BROWSER_SLOTS` задаётся в `.env` (по умолчанию 10 — на ~10–12 ядер). -- Сверх `BROWSER_SLOTS` запросы сразу получают `503` с `code: no_slots`, не ставятся в очередь. +- Правило подбора: `BROWSER_SLOTS ≈ число ядер`, `WORKER_MEM_LIMIT ≈ 0,6 ГБ × BROWSER_SLOTS + 1 ГБ`. + Пример: сервер 4 ядра / 8 ГБ → 3–4 слота. `BROWSER_SLOTS` задаётся в `.env`. +- Удерживаемые окна (`browser_idle_timeout`) занимают слоты, но почти не тратят CPU. +- Сверх `BROWSER_SLOTS` задачи v2 ждут в очереди, а `/v1/solve` сразу получает `503` с `code: no_slots`. + +Замеры выше сделаны на v1 (один сервис, Chrome в api); схема нагрузки на Chrome с очередью та же. Повторить замер: @@ -104,4 +132,4 @@ api.example.com { } ``` -Таймауты прокси должны быть не меньше `timeout` запроса (по умолчанию 120 с). +Таймауты прокси для `/v1/solve` должны быть не меньше `timeout` запроса (по умолчанию 120 с); запросы v2 короткие. diff --git a/docs/Usage.md b/docs/Usage.md index a7b5d19..db9e165 100644 --- a/docs/Usage.md +++ b/docs/Usage.md @@ -46,14 +46,29 @@ UA в ответе может измениться, когда на сервер localStorage, отпечаток), и повторно проходить антибот обычно не нужно — запрос заметно быстрее. - Одна сессия — один «человек»: держите для неё **один и тот же прокси**. Смена exit-IP внутри сессии выглядит подозрительно и обнуляет `cf_clearance`. -- Не отправляйте два запроса одной сессии одновременно: профиль Chrome занят первым запросом. +- Задачи одной сессии выполняются по очереди: вторая ждёт, пока закончится первая. - Для параллельной работы заводите несколько сессий, каждую со своим прокси. +## Синхронно или через задачи + +- `/v1/solve` — проще: один запрос, один ответ. Держите соединение открытым до 3–4 минут (см. «Таймауты»). +- `/v2/tasks` — запрос возвращается сразу, результат забираете по ссылке `GET /v2/tasks/{id}` с той же cookie + `session_id`. Подходит, когда соединение нельзя держать долго, и для многошаговых сценариев. + +Многошаговый сценарий (войти, затем несколько раз что-то сделать на сайте) — с удержанием окна: + +1. `POST /v2/tasks` с `url` и `browser_idle_timeout`: окно останется открытым на это число минут простоя. +2. Следующие задачи с той же cookie `session_id` и тем же `proxy` выполняются в открытом окне: без запуска Chrome + и повторного антибота. Без `url` — на текущей странице. +3. В конце — `DELETE /v2/sessions/current` или задача с `browser_idle_timeout: 0`: окно закроется и освободит слот. + +Открытое окно занимает слот браузера, поэтому ставьте простой не больше, чем пауза между вашими шагами. + ## Cookies - Переданные в запросе `cookies` **заменяют** cookies профиля целиком. Не передали — профиль сохраняет свои. -- Для постоянных cookies указывайте `expires` (unix-время). Cookie без `expires` — сессионная и пропадёт между - запросами, потому что браузер перезапускается на каждый запрос. +- Для постоянных cookies указывайте `expires` (unix-время). Cookie без `expires` — сессионная и пропадёт, как + только закроется окно браузера (после задачи без `browser_idle_timeout`). - Обычно передавать cookies не нужно: профиль сам хранит то, что выдал сайт. Передавайте, только если сайт ведёт себя по-разному в зависимости от cookies (например, авторизация, полученная вне Robotex). @@ -61,7 +76,7 @@ UA в ответе может измениться, когда на сервер | Ответ | Что делать | |---|---| -| `503`, `code: no_slots` | все браузеры заняты — повторить через `Retry-After` секунд | +| `503`, `code: no_slots` | `/v1/solve`: все браузеры заняты — повторить через `Retry-After` секунд. Задачи v2 в этом случае ждут в очереди | | `challenge_failed` | антибот не пройден за отведённое время: повторить в той же сессии; при повторных неудачах — сменить прокси (и сессию) | | `action_failed` | упало действие: смотрите `error` и `actions[].result`; для отладки повторите с `"screenshot": true` | | `request_failed` | навигация не удалась: обычно прокси недоступен или сайт не отвечает | diff --git a/dot_env b/dot_env index 8713326..8a2efdf 100644 --- a/dot_env +++ b/dot_env @@ -1,4 +1,5 @@ API_PORT= -API_CPUS=1 -API_MEM_LIMIT=1g +WORKER_CPUS=1 +WORKER_MEM_LIMIT=1g BROWSER_SLOTS=1 +RABBITMQ_PASSWORD= diff --git a/env.example b/env.example index 78de7d9..4689048 100644 --- a/env.example +++ b/env.example @@ -1,9 +1,16 @@ # Project Env -# порт, CPU и память контейнера (docker-compose) +# порт api на хосте (docker-compose) API_PORT=8000 -API_CPUS=4 -API_MEM_LIMIT=8g -# сколько Chrome параллельно, ≈ числу ядер (docs/Deployment.md) +# CPU и память контейнера worker +WORKER_CPUS=4 +WORKER_MEM_LIMIT=8g +# сколько окон Chrome на воркер — и выполняющих задачу, и удерживаемых (browser_idle_timeout); ≈ числу ядер BROWSER_SLOTS=4 +# пароль RabbitMQ (сеть compose наружу не открыта) +RABBITMQ_PASSWORD=robotex +# сколько часов хранится результат задачи в Redis +RESULT_TTL_HOURS=24 +# сколько /v1/solve ждёт, пока воркер возьмёт задачу, прежде чем ответить 503 +V1_QUEUE_WAIT=30 NO_SLOTS_RETRY_AFTER=5 LOG_LEVEL=INFO diff --git a/requirements-api.txt b/requirements-api.txt new file mode 100644 index 0000000..0426c45 --- /dev/null +++ b/requirements-api.txt @@ -0,0 +1,17 @@ +# зависимости образа api (Dockerfile.api): без Chrome и Patchright. Версии — как в requirements.txt +annotated-doc==0.0.5 +annotated-types==0.8.0 +anyio==4.15.1 +click==8.5.0 +dramatiq==2.2.1 +fastapi==0.141.1 +h11==0.16.0 +idna==3.19 +pika==1.4.4 +pydantic==2.13.5 +pydantic_core==2.46.5 +redis==8.1.0 +starlette==1.6.0 +typing-inspection==0.4.4 +typing_extensions==4.16.0 +uvicorn==0.52.4 diff --git a/requirements-dev.txt b/requirements-dev.txt index c341c30..1d8f85a 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -3,3 +3,5 @@ pytest==9.1.1 pluggy==1.6.0 iniconfig==2.3.1 packaging==26.3 +fakeredis==2.39.0 +sortedcontainers==2.4.0 diff --git a/requirements.txt b/requirements.txt index 3031a81..de459fe 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,6 +15,7 @@ certifi==2026.7.22 charset-normalizer==3.5.1 click==8.5.0 distro==1.9.0 +dramatiq==2.2.1 editor==1.8.0 fastapi==0.141.1 frozenlist==1.8.0 @@ -47,6 +48,7 @@ multidict==6.7.1 numpy==2.5.3 orjson==3.12.0 patchright==1.63.0 +pika==1.4.4 platformdirs==4.11.7 playwright==1.62.0 propcache==0.5.2 diff --git a/src/api/logs.py b/src/api/logs.py new file mode 100644 index 0000000..752d2b7 --- /dev/null +++ b/src/api/logs.py @@ -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) diff --git a/src/api/schemas.py b/src/api/schemas.py index def9752..a8780b8 100644 --- a/src/api/schemas.py +++ b/src/api/schemas.py @@ -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 diff --git a/src/engine/stealthy.py b/src/engine/stealthy.py index cd8d051..3ed19ad 100644 --- a/src/engine/stealthy.py +++ b/src/engine/stealthy.py @@ -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 diff --git a/src/main.py b/src/main.py index 3658b26..8ea69f1 100644 --- a/src/main.py +++ b/src/main.py @@ -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 "") + 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 diff --git a/src/tasks/__init__.py b/src/tasks/__init__.py new file mode 100644 index 0000000..5528462 --- /dev/null +++ b/src/tasks/__init__.py @@ -0,0 +1,5 @@ +""" +Общее для сервисов api и worker (#78): хранилище задач в Redis, брокер dramatiq, сессии. + +Пакет не зависит от браузера: образ api собирается без Chrome и Patchright. +""" diff --git a/src/tasks/broker.py b/src/tasks/broker.py new file mode 100644 index 0000000..e584878 --- /dev/null +++ b/src/tasks/broker.py @@ -0,0 +1,52 @@ +""" +Брокер dramatiq поверх RabbitMQ (#78) и сообщения, которыми api ставит работу воркерам. + +Каждый воркер слушает общую очередь `robotex` и персональную `robotex.worker.`. В персональную +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={}) diff --git a/src/tasks/sessions.py b/src/tasks/sessions.py new file mode 100644 index 0000000..78d1324 --- /dev/null +++ b/src/tasks/sessions.py @@ -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()) diff --git a/src/tasks/store.py b/src/tasks/store.py new file mode 100644 index 0000000..4299eb2 --- /dev/null +++ b/src/tasks/store.py @@ -0,0 +1,199 @@ +""" +Состояние задач, окон и воркеров в Redis (#78). Сроки жизни — TTL ключей Redis. + +* `robotex:task:` (hash) — статус задачи, сессия, время, результат. Пока задача не завершена, + у ключа страховочный TTL `PENDING_TTL`; при завершении — `RESULT_TTL` (`RESULT_TTL_HOURS`). +* `robotex:session:` (hash) — где открыто окно сессии: `worker_id` и `idle_timeout`. + Пока задача выполняется, TTL — `RUNNING_TTL`; после неё — `browser_idle_timeout`. Ключ исчез — + окно закрывается (worker/windows.py). +* `robotex:worker:` — 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"]) diff --git a/src/worker/__init__.py b/src/worker/__init__.py new file mode 100644 index 0000000..9d71285 --- /dev/null +++ b/src/worker/__init__.py @@ -0,0 +1,6 @@ +""" +Сервис worker (#78): берёт задачи из очереди RabbitMQ, выполняет их в Chrome и кладёт результат в Redis. + +Запуск: `dramatiq worker.actors --processes 1 --threads N` под Xvfb (см. Dockerfile). Процесс — строго +один: открытые окна браузера живут в его памяти. +""" diff --git a/src/worker/actors.py b/src/worker/actors.py new file mode 100644 index 0000000..779e170 --- /dev/null +++ b/src/worker/actors.py @@ -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)) diff --git a/src/worker/windows.py b/src/worker/windows.py new file mode 100644 index 0000000..eca02e8 --- /dev/null +++ b/src/worker/windows.py @@ -0,0 +1,348 @@ +""" +Слоты браузера и окна сессий воркера (#81, #82). + +Слот — один открытый Chrome. `BROWSER_SLOTS` ограничивает число окон воркера: и тех, где идёт задача, +и удерживаемых между задачами. Окно сессии остаётся открытым после задачи, если в ней задан +`browser_idle_timeout > 0`: следующая задача с той же cookie `session_id` попадает в это окно +(маршрутизация — в api, #83) и продолжает работу без перезапуска Chrome. Время простоя считает +Redis: TTL ключа `session:`, каждая задача в окно ставит его заново. Ключ исчез — окно +закрывается фоновой проверкой (`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 "") + 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) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/browser/__init__.py b/tests/browser/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/browser/test_e2e.py b/tests/browser/test_e2e.py index 6f85e98..b4b890c 100644 --- a/tests/browser/test_e2e.py +++ b/tests/browser/test_e2e.py @@ -1,16 +1,19 @@ -"""Сквозные тесты с настоящим Chrome и сетью. Запускаются в Docker (нужны Chrome и Xvfb): +"""Сквозные тесты с настоящим Chrome и сетью: api и worker в одном процессе (tests/harness.py), без +RabbitMQ и Redis. Запускаются в Docker (нужны Chrome и Xvfb): make test-browser + +Весь стенд с очередью проверяет make smoke. """ import json import re import uuid import pytest -from fastapi.testclient import TestClient -import main from api import profile +from tests.harness import run_env +from worker import windows pytestmark = pytest.mark.browser @@ -34,9 +37,14 @@ FP_JS = """async () => { @pytest.fixture -def client(monkeypatch, tmp_path): - monkeypatch.setattr(main, "USER_DATA_ROOT", tmp_path) - return TestClient(main.app) +def env(monkeypatch, tmp_path): + with run_env(monkeypatch, tmp_path) as harness: + yield harness + + +@pytest.fixture +def client(env): + return env.client def solve(client, session_id=None, **body): @@ -50,11 +58,11 @@ def solve(client, session_id=None, **body): def test_basic_page(client): data = solve(client, url="https://example.com", screenshot=True, actions=[ - {"type": "wait_for", "selector": "a", "timeout": 10}, - {"type": "inner_html", "selector": "a"}, + {"type": "wait_for", "selector": "p", "timeout": 10}, + {"type": "inner_html", "selector": "title"}, ]) assert data["status"] == "ok", data - assert data["actions"][1]["result"] == "Learn more" + assert data["actions"][1]["result"] == "Example Domain" assert len(data["screenshot"]) > 1000 @@ -88,7 +96,7 @@ def test_fingerprint_applied_and_stable(client, monkeypatch): return json.loads(await page.inner_html("#fp")) results = {} - original = main.make_page_action + original = windows.make_page_action def capture(payload): action = original(payload) @@ -99,22 +107,22 @@ def test_fingerprint_applied_and_stable(client, monkeypatch): return out return wrapped - monkeypatch.setattr(main, "make_page_action", capture) + monkeypatch.setattr(windows, "make_page_action", capture) - class Spy(main.AsyncStealthySession): - async def fetch(self, url, page_action): - out = await super().fetch(url, page_action=page_action) + class Spy(windows.AsyncStealthySession): + async def fetch(self, url, page_action, **kwargs): + out = await super().fetch(url, page_action=page_action, **kwargs) results.setdefault(self._user_data_dir, []).append(out["fp"]) return out - monkeypatch.setattr(main, "AsyncStealthySession", Spy) + monkeypatch.setattr(windows, "AsyncStealthySession", Spy) sids = [solve(client, url="https://example.com")["session_id"] for _ in range(4)] solve(client, sids[0], url="https://example.com") fps = {} for d, measured in results.items(): - saved = profile.load(main.USER_DATA_ROOT / d.rsplit("/", 1)[-1]) + saved = profile.load(windows.USER_DATA_ROOT / d.rsplit("/", 1)[-1]) fp = saved["fingerprint"] for m in measured: assert m["screen"] == [fp["screen"]["width"], fp["screen"]["height"], fp["device_scale_factor"]] @@ -141,3 +149,24 @@ def test_cloudflare(client): assert data["status"] == "ok", data.get("error") assert "cf_clearance" in {c["name"] for c in data["cookies"]} assert uuid.UUID(data["session_id"]) + + +def test_window_held_between_tasks(env): + """browser_idle_timeout: вторая задача без url выполняется в том же окне, на уже открытой странице.""" + read = [{"type": "inner_html", "selector": "title"}] + client = env.client + + first = client.post("/v2/tasks", json={"url": "https://example.com", "actions": read, "browser_idle_timeout": 1}) + task = client.get(first.json()["url"]).json() + assert task["result"]["status"] == "ok", task + assert env.pool.used == 1 + + second = client.post("/v2/tasks", json={"actions": read, "browser_idle_timeout": 1}).json() + result = client.get(second["url"]).json()["result"] + assert result["status"] == "ok", result + assert result["actions"][0]["result"] == "Example Domain" + # окно не перезапускалось: страница не открывалась заново + assert result["duration_ms"] < task["result"]["duration_ms"] + + assert client.delete("/v2/sessions/current").json() == {"closed": True} + assert env.pool.used == 0 diff --git a/tests/conftest.py b/tests/conftest.py index 2c61d83..3e2d5cd 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,10 @@ +import os + import pytest +# брокер dramatiq в памяти: tasks.broker создаётся при импорте +os.environ.setdefault("RABBITMQ_URL", "stub://") + @pytest.fixture def anyio_backend(): diff --git a/tests/harness.py b/tests/harness.py new file mode 100644 index 0000000..395d707 --- /dev/null +++ b/tests/harness.py @@ -0,0 +1,65 @@ +"""api и worker в одном процессе без RabbitMQ и Redis: Redis — fakeredis, отправка в очередь подменена +доставкой сообщения прямо в WindowPool воркера (tests/unit, tests/browser).""" +import json +from contextlib import contextmanager + +import fakeredis +from fastapi.testclient import TestClient + +import main +from tasks import store as task_store +from tasks.broker import RUN_TASK +from worker import windows + +WORKER_ID = "w1" + + +class Harness: + def __init__(self, store, pool, redis_sync): + self.store = store + self.pool = pool + self.redis = redis_sync + self.messages = [] + self.outcomes = [] + # True — сообщение выполняется сразу при отправке; False — копится в messages до deliver() + self.auto = True + self.client: TestClient | None = None + + async def dispatch(self, message): + self.messages.append(message) + if self.auto: + await self.deliver(message) + + async def deliver(self, message): + if message.actor_name == RUN_TASK: + self.outcomes.append(await self.pool.handle(*message.args)) + else: + await self.pool.close(*message.args) + await self.store.heartbeat(WORKER_ID, self.pool.info()) + + def deliver_pending(self): + pending, self.messages = self.messages, [] + for message in pending: + self.client.portal.call(self.deliver, message) + + def heartbeat(self, **info): + self.redis.set(f"{task_store.PREFIX}:worker:{WORKER_ID}", json.dumps(self.pool.info() | info)) + + def call(self, fn, *args): + return self.client.portal.call(fn, *args) + + +@contextmanager +def run_env(monkeypatch, tmp_path, slots=2): + server = fakeredis.FakeServer() + store = task_store.TaskStore(fakeredis.aioredis.FakeRedis(server=server, decode_responses=True)) + pool = windows.WindowPool(store, WORKER_ID, slots=slots) + harness = Harness(store, pool, fakeredis.FakeRedis(server=server, decode_responses=True)) + monkeypatch.setattr(windows, "USER_DATA_ROOT", tmp_path) + monkeypatch.setattr(main, "store", store) + monkeypatch.setattr(main, "dispatch", harness.dispatch) + harness.heartbeat() + with TestClient(main.app) as client: + harness.client = client + yield harness + client.portal.call(pool.close_all) diff --git a/tests/smoke.py b/tests/smoke.py new file mode 100644 index 0000000..341f741 --- /dev/null +++ b/tests/smoke.py @@ -0,0 +1,76 @@ +"""Проверка запущенного стенда (make up): api, очередь, воркер с Chrome, удержание окна. + + make smoke # http://localhost:8000 + make smoke API_URL=https://app.robotex.sokolab.xyz +""" +import sys +import time + +import httpx + +URL = "https://example.com" +READ_TITLE = [{"type": "inner_html", "selector": "title"}] + + +def check(cond: bool, what: str) -> None: + print((" ok " if cond else " FAIL ") + what) + if not cond: + sys.exit(1) + + +def wait_task(client: httpx.Client, task_id: str, timeout: float = 180) -> dict: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + task = client.get(f"/v2/tasks/{task_id}").json() + if task["status"] not in ("queued", "running"): + return task + time.sleep(0.5) + raise TimeoutError(f"task {task_id} not finished in {timeout}s") + + +def windows(client: httpx.Client) -> int: + return sum(w["windows"] for w in client.get("/health").json()["workers"]) + + +def main(base_url: str) -> None: + client = httpx.Client(base_url=base_url, timeout=300) + + print("health") + health = client.get("/health").json() + check(len(health["workers"]) > 0, f"workers: {[w['worker_id'] for w in health['workers']]}, free slots {health['free_slots']}") + + print("v1 /v1/solve") + started = time.monotonic() + data = client.post("/v1/solve", json={"url": URL, "actions": READ_TITLE}).json() + check(data["status"] == "ok", f"status {data['status']} {data.get('error', '')} ({time.monotonic() - started:.1f}s)") + check(data["actions"][0]["result"] == "Example Domain", "title = Example Domain") + + print("v2 window held between tasks") + v2 = httpx.Client(base_url=base_url, timeout=60) # своя сессия + started = time.monotonic() + created = v2.post("/v2/tasks", json={"url": URL, "actions": READ_TITLE, "browser_idle_timeout": 2}) + check(created.status_code == 202, f"POST /v2/tasks -> {created.status_code}") + task = wait_task(v2, created.json()["task_id"]) + check(task["status"] == "done" and task["result"]["status"] == "ok", f"first task done ({time.monotonic() - started:.1f}s)") + check(windows(client) >= 1, "window stays open") + + started = time.monotonic() + created = v2.post("/v2/tasks", json={"actions": READ_TITLE, "browser_idle_timeout": 2}) + task = wait_task(v2, created.json()["task_id"]) + check( + task["status"] == "done" and task["result"]["actions"][0]["result"] == "Example Domain", + f"second task without url ran in the open window ({time.monotonic() - started:.1f}s)", + ) + + check(client.get(f"/v2/tasks/{created.json()['task_id']}").status_code == 404, "task hidden from other sessions") + + check(v2.delete("/v2/sessions/current").json() == {"closed": True}, "DELETE /v2/sessions/current") + deadline = time.monotonic() + 15 + while windows(client) and time.monotonic() < deadline: + time.sleep(0.5) + check(windows(client) == 0, "window closed, slot released") + print("all good") + + +if __name__ == "__main__": + main(sys.argv[1] if len(sys.argv) > 1 else "http://localhost:8000") diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 0000000..3377a0f --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,80 @@ +"""Сервисы api и worker без браузера: Chrome — FakeSession (окружение — tests/harness.py).""" +import pytest + +from api import actions +from tests.harness import WORKER_ID, run_env # noqa: F401 +from worker import windows + + +class FakeContext: + def __init__(self): + self.ops = [] + + async def cookies(self): + return [{"name": "cf_clearance", "value": "secret-cookie"}] + + async def clear_cookies(self): + self.ops.append("clear") + + async def add_cookies(self, cookies): + self.ops.append(("add", cookies)) + + +class FakePage: + def __init__(self, context): + self.context = context + + async def screenshot(self, type): + return b"png" + + async def inner_html(self, selector): + return "page" + + async def evaluate(self, js): + return "UA/engine" + + +class FakeSession: + """Вместо AsyncStealthySession: запоминает kwargs и вызовы, page_action выполняет на FakePage.""" + instances: list["FakeSession"] = [] + error: Exception | None = None + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.context = FakeContext() + self.calls = [] + self.closed = False + FakeSession.instances.append(self) + + async def start(self): + pass + + async def close(self): + self.closed = True + + async def fetch(self, url, page_action, **kwargs): + self.calls.append(("fetch", url)) + if FakeSession.error: + raise FakeSession.error + return await page_action(FakePage(self.context)) + + async def run(self, page_action, **kwargs): + self.calls.append(("run", None)) + return await page_action(FakePage(self.context)) + + +def fake_execute(result=None, exc=None): + async def execute(*args, **kwargs): + if exc: + raise exc + return result if result is not None else [] + return execute + + +@pytest.fixture +def env(monkeypatch, tmp_path): + FakeSession.instances, FakeSession.error = [], None + monkeypatch.setattr(windows, "AsyncStealthySession", FakeSession) + monkeypatch.setattr(actions, "execute", fake_execute()) + with run_env(monkeypatch, tmp_path) as harness: + yield harness diff --git a/tests/unit/test_actors.py b/tests/unit/test_actors.py new file mode 100644 index 0000000..9646810 --- /dev/null +++ b/tests/unit/test_actors.py @@ -0,0 +1,63 @@ +"""Настоящий воркер dramatiq на брокере в памяти: actor'ы, event loop AsyncIO, пул окон и heartbeat.""" +import asyncio + +import dramatiq +import fakeredis +import pytest + +from api import actions +from tasks import store as task_store +from tasks.broker import COMMON_QUEUE, personal_queue, task_message +from worker import windows + +from .conftest import FakeSession, fake_execute + + +@pytest.fixture +def worker(monkeypatch, tmp_path): + server = fakeredis.FakeServer() + monkeypatch.setattr(task_store, "connect", lambda url=None: fakeredis.aioredis.FakeRedis(server=server, decode_responses=True)) + monkeypatch.setattr(windows, "AsyncStealthySession", FakeSession) + monkeypatch.setattr(windows, "USER_DATA_ROOT", tmp_path) + monkeypatch.setattr(actions, "execute", fake_execute()) + FakeSession.instances, FakeSession.error = [], None + + from worker import actors + w = dramatiq.Worker(actors.broker, worker_timeout=100, worker_threads=2) + w.start() + try: + yield actors, server + finally: + w.stop() + # при остановке воркер закрывает окна и снимает heartbeat + assert all(s.closed for s in FakeSession.instances) + assert not fakeredis.FakeRedis(server=server).keys("robotex:worker:*") + + +SID = "11111111-1111-1111-1111-111111111111" + + +def submit(actors, server, queue, idle=0): + """То же, что делает api: задача в Redis, сообщение в очередь; ждём, пока воркер её выполнит.""" + store = task_store.TaskStore(fakeredis.aioredis.FakeRedis(server=server, decode_responses=True)) + asyncio.run(store.create("t1", SID, "https://example.com")) + actors.broker.enqueue(task_message("t1", SID, {"url": "https://example.com", "actions": []}, idle, queue)) + actors.broker.join(queue) + return fakeredis.FakeRedis(server=server, decode_responses=True) + + +@pytest.mark.parametrize("queue", ["common", "personal"]) +def test_task_executed_from_queue(worker, queue): + actors, server = worker + redis = submit(actors, server, COMMON_QUEUE if queue == "common" else personal_queue(actors.WORKER_ID)) + task = redis.hgetall("robotex:task:t1") + assert task["status"] == "done", task + assert '"status": "ok"' in task["result"] + assert redis.exists(f"robotex:worker:{actors.WORKER_ID}") + + +def test_window_held_until_shutdown(worker): + actors, server = worker + redis = submit(actors, server, COMMON_QUEUE, idle=5) + assert redis.hgetall(f"robotex:session:{SID}")["worker_id"] == actors.WORKER_ID + assert not FakeSession.instances[0].closed diff --git a/tests/unit/test_solve.py b/tests/unit/test_solve.py deleted file mode 100644 index e66ccec..0000000 --- a/tests/unit/test_solve.py +++ /dev/null @@ -1,240 +0,0 @@ -"""/v1/solve без браузера: движок подменён, проверяются сборка ответа, коды ошибок, слоты, сессии и логи.""" -import asyncio -import logging -import uuid - -import pytest -from fastapi.testclient import TestClient - -import main -from api import actions, profile - - -class FakeContext: - async def cookies(self): - return [{"name": "cf_clearance", "value": "secret-cookie"}] - - -class FakePage: - context = FakeContext() - - async def screenshot(self, type): - return b"png" - - async def inner_html(self, selector): - return "page" - - async def evaluate(self, js): - return "UA/engine" - - -class FakeSession: - """Вместо AsyncStealthySession: запоминает kwargs и вызывает page_action на FakePage.""" - calls: list[dict] = [] - error: Exception | None = None - - def __init__(self, **kwargs): - FakeSession.calls.append(kwargs) - - async def __aenter__(self): - return self - - async def __aexit__(self, *exc): - return False - - async def fetch(self, url, page_action): - if FakeSession.error: - raise FakeSession.error - return await page_action(FakePage()) - - -@pytest.fixture -def client(monkeypatch, tmp_path): - FakeSession.calls, FakeSession.error = [], None - monkeypatch.setattr(main, "AsyncStealthySession", FakeSession) - monkeypatch.setattr(main, "USER_DATA_ROOT", tmp_path) - monkeypatch.setattr(main, "_slots", asyncio.Semaphore(2)) - return TestClient(main.app) - - -def fake_execute(result=None, exc=None): - async def execute(*args, **kwargs): - if exc: - raise exc - return result - return execute - - -def post(client, **body): - return client.post("/v1/solve", json={"url": "https://example.com", "actions": [], **body}) - - -def test_ok(client, monkeypatch): - monkeypatch.setattr(actions, "execute", fake_execute([{"type": "wait_for", "ok": True, "duration_ms": 5}])) - resp = post(client, html=True, screenshot=True) - - assert resp.status_code == 200 - data = resp.json() - assert data["status"] == "ok" - assert "code" not in data and "error" not in data - assert data["user_agent"] == "UA/engine" - assert data["actions"] == [{"type": "wait_for", "ok": True, "duration_ms": 5}] - assert data["html"] == "page" - assert data["screenshot"] == "cG5n" # base64(b"png") - assert data["duration_ms"] >= 0 - assert uuid.UUID(resp.cookies["session_id"]) - - -def test_html_and_screenshot_only_on_request(client, monkeypatch): - monkeypatch.setattr(actions, "execute", fake_execute([])) - data = post(client).json() - assert "html" not in data and "screenshot" not in data - assert data["actions"] == [] - - -def test_action_failed(client, monkeypatch): - monkeypatch.setattr(actions, "execute", fake_execute([ - {"type": "click", "ok": True, "duration_ms": 1}, - {"type": "inner_html", "ok": False, "result": "selector not found: #x", "duration_ms": 1}, - ])) - data = post(client, screenshot=True).json() - assert data["status"] == "err" - assert data["code"] == "action_failed" - assert data["error"] == "inner_html: selector not found: #x" - assert data["screenshot"] - - -def test_challenge_failed(client, monkeypatch): - monkeypatch.setattr(actions, "execute", fake_execute(exc=actions.ChallengeFailed("cloudflare not passed: timeout"))) - data = post(client, screenshot=True).json() - assert data["status"] == "err" - assert data["code"] == "challenge_failed" - assert data["error"] == "cloudflare not passed: timeout" - assert data["actions"] == [] - assert data["cookies"] and data["screenshot"] - - -def test_request_failed(client): - FakeSession.error = RuntimeError("net::ERR_PROXY_CONNECTION_FAILED") - data = post(client).json() - assert data == { - "status": "err", "code": "request_failed", "error": "net::ERR_PROXY_CONNECTION_FAILED", - "actions": [], "duration_ms": data["duration_ms"], - } - - -def test_no_slots(client, monkeypatch): - monkeypatch.setattr(main, "_slots", asyncio.Semaphore(0)) - resp = post(client) - assert resp.status_code == 503 - assert resp.headers["retry-after"] == main.RETRY_AFTER - assert resp.json() == {"status": "err", "code": "no_slots", "error": "No slots available"} - assert FakeSession.calls == [] - - -def test_slot_released_after_request(client, monkeypatch): - sem = asyncio.Semaphore(1) - monkeypatch.setattr(main, "_slots", sem) - monkeypatch.setattr(actions, "execute", fake_execute([])) - FakeSession.error = RuntimeError("boom") - for _ in range(3): - assert post(client).status_code == 200 - assert not sem.locked() - - -def test_invalid_request_is_422_without_browser(client): - resp = post(client, actions=[{"type": "nope"}]) - assert resp.status_code == 422 - assert FakeSession.calls == [] - - -def test_session_cookie_reused(client, monkeypatch, tmp_path): - monkeypatch.setattr(actions, "execute", fake_execute([])) - sid = str(uuid.uuid4()) - client.cookies.set("session_id", sid) - resp = post(client) - assert resp.cookies["session_id"] == sid - assert FakeSession.calls[0]["user_data_dir"] == str(tmp_path / sid) - - -@pytest.mark.parametrize("bad", ["../../etc", "not-a-uuid", "/tmp/x"]) -def test_bad_session_cookie_gets_new_session(client, monkeypatch, tmp_path, bad): - monkeypatch.setattr(actions, "execute", fake_execute([])) - client.cookies.set("session_id", bad) - resp = post(client) - sid = resp.cookies["session_id"] - assert sid != bad and uuid.UUID(sid) - assert [p.name for p in tmp_path.iterdir()] == [sid] - - -def test_user_agent_param_rejected(client): - resp = post(client, user_agent="First/1") - assert resp.status_code == 422 - assert "take user_agent from the response" in resp.text - assert FakeSession.calls == [] - - -def test_fingerprint_passed_to_browser_and_ua_untouched(client, monkeypatch): - monkeypatch.setattr(actions, "execute", fake_execute([])) - post(client) - kwargs = FakeSession.calls[0] - assert "useragent" not in kwargs - assert set(kwargs["additional_args"]) == {"screen", "device_scale_factor"} - - -def test_cookies_replace_profile_cookies_only_when_passed(client, monkeypatch): - monkeypatch.setattr(actions, "execute", fake_execute([])) - post(client) - post(client, cookies=[{"name": "a", "value": "1", "url": "https://example.com"}]) - post(client, cookies=[]) - assert [c["clear_cookies"] for c in FakeSession.calls] == [False, True, True] - - -def test_proxy_passed_and_geo_resolved(client, monkeypatch): - monkeypatch.setattr(actions, "execute", fake_execute([])) - seen = {} - - async def fake_geo(proxy): - seen["proxy"] = proxy - return {"locale": "de-DE", "timezone_id": "Europe/Berlin"} - - monkeypatch.setattr(main, "resolve_geo", fake_geo) - post(client, proxy="http://u:p@host:3128") - kwargs = FakeSession.calls[0] - assert kwargs["proxy"] == {"server": "http://host:3128", "username": "u", "password": "p"} - assert seen["proxy"] == kwargs["proxy"] - assert (kwargs["locale"], kwargs["timezone_id"]) == ("de-DE", "Europe/Berlin") - - -def test_without_proxy_no_geo(client, monkeypatch): - monkeypatch.setattr(actions, "execute", fake_execute([])) - - async def fail_geo(proxy): - raise AssertionError("geo must not be resolved without proxy") - - monkeypatch.setattr(main, "resolve_geo", fail_geo) - post(client) - assert FakeSession.calls[0]["proxy"] is None - - -def test_secrets_not_logged(client, monkeypatch, caplog): - monkeypatch.setattr(actions, "execute", fake_execute([])) - - async def no_geo(proxy): - return None - - monkeypatch.setattr(main, "resolve_geo", no_geo) - caplog.set_level(logging.DEBUG, logger="main") - post( - client, proxy="http://someuser:TopSecret1@host:3128", captcha="KEYSECRET2", html=True, - cookies=[{"name": "sid", "value": "CookieSecret3", "url": "https://example.com"}], - actions=[{"type": "fill", "selector": "#p", "value": "PassSecret4"}], - ) - log = caplog.text - assert "Incoming payload" in log and "Outgoing payload" in log - for secret in ("someuser", "TopSecret1", "KEYSECRET2", "CookieSecret3", "PassSecret4", "secret-cookie", "page"): - assert secret not in log - - -def test_health(client): - assert client.get("/health").json() == {"status": "ok", "browser_slots": main.BROWSER_SLOTS, "free_slots": 2} diff --git a/tests/unit/test_store.py b/tests/unit/test_store.py new file mode 100644 index 0000000..33e2336 --- /dev/null +++ b/tests/unit/test_store.py @@ -0,0 +1,74 @@ +import fakeredis +import pytest + +from tasks import store as task_store +from tasks.store import Status, TaskStore + +pytestmark = pytest.mark.anyio + + +@pytest.fixture +def store(): + return TaskStore(fakeredis.aioredis.FakeRedis(decode_responses=True)) + + +async def test_claim_once(store): + await store.create("t", "s", "https://example.com") + await store.heartbeat("w1", {}) + assert (await store.get("t"))["status"] == Status.QUEUED + assert await store.claim("t", "w1") + assert not await store.claim("t", "w2") + assert not await store.cancel("t") + + +async def test_cancelled_task_is_not_claimed(store): + await store.create("t", "s", None) + assert await store.cancel("t") + assert not await store.claim("t", "w1") + assert (await store.get("t"))["status"] == Status.CANCELLED + + +async def test_missing_task_is_not_claimed(store): + assert not await store.claim("nope", "w1") + + +async def test_redelivery_while_running_here_is_skipped(store): + await store.create("t", "s", None) + await store.heartbeat("w1", {}) + assert await store.claim("t", "w1") + assert not await store.claim("t", "w1", running_here=True) + + +async def test_task_of_dead_worker_is_reclaimed_limited_times(store): + await store.create("t", "s", None) + await store.heartbeat("w2", {}) + assert await store.claim("t", "w1") # w1 умер — heartbeat нет + assert await store.claim("t", "w2") # attempts = 2 + await store.remove_worker("w2") + assert not await store.claim("t", "w3") + task = await store.get("t") + assert task["status"] == Status.FAILED and "worker died" in task["error"] + + +async def test_task_of_live_worker_is_not_reclaimed(store): + await store.create("t", "s", None) + await store.heartbeat("w1", {}) + assert await store.claim("t", "w1") + assert not await store.claim("t", "w2") + + +async def test_finish_sets_result_and_ttl(store): + await store.create("t", "s", None) + assert await store.redis.ttl("robotex:task:t") == task_store.PENDING_TTL + await store.finish("t", {"status": "ok", "actions": []}) + task = await store.get("t") + assert task["status"] == Status.DONE and task["result"] == {"status": "ok", "actions": []} + assert await store.redis.ttl("robotex:task:t") == task_store.RESULT_TTL + + +async def test_live_window_requires_live_worker(store): + await store.set_session("s", "w1", 5, 300) + assert await store.live_window("s") is None + await store.heartbeat("w1", {"browser_slots": 1, "free_slots": 0}) + assert await store.live_window("s") == {"worker_id": "w1", "idle_timeout": "5"} + assert await store.workers() == [{"worker_id": "w1", "browser_slots": 1, "free_slots": 0}] diff --git a/tests/unit/test_v1.py b/tests/unit/test_v1.py new file mode 100644 index 0000000..7c2bd8b --- /dev/null +++ b/tests/unit/test_v1.py @@ -0,0 +1,200 @@ +"""/v1/solve поверх очереди: контракт v1 не изменился (#84).""" +import logging +import uuid + +import pytest + +import main +from api import actions +from worker import windows +from worker.windows import Outcome + +from .conftest import FakeSession, fake_execute + + +def post(env, **body): + return env.client.post("/v1/solve", json={"url": "https://example.com", "actions": [], **body}) + + +def test_ok(env, monkeypatch): + monkeypatch.setattr(actions, "execute", fake_execute([{"type": "wait_for", "ok": True, "duration_ms": 5}])) + resp = post(env, html=True, screenshot=True) + + assert resp.status_code == 200 + data = resp.json() + assert data["status"] == "ok" + assert "code" not in data and "error" not in data + assert data["user_agent"] == "UA/engine" + assert data["actions"] == [{"type": "wait_for", "ok": True, "duration_ms": 5}] + assert data["html"] == "page" + assert data["screenshot"] == "cG5n" # base64(b"png") + assert data["duration_ms"] >= 0 + assert uuid.UUID(resp.cookies["session_id"]) + + +def test_html_and_screenshot_only_on_request(env): + data = post(env).json() + assert "html" not in data and "screenshot" not in data + assert data["actions"] == [] + + +def test_action_failed(env, monkeypatch): + monkeypatch.setattr(actions, "execute", fake_execute([ + {"type": "click", "ok": True, "duration_ms": 1}, + {"type": "inner_html", "ok": False, "result": "selector not found: #x", "duration_ms": 1}, + ])) + data = post(env, screenshot=True).json() + assert data["status"] == "err" + assert data["code"] == "action_failed" + assert data["error"] == "inner_html: selector not found: #x" + assert data["screenshot"] + + +def test_challenge_failed(env, monkeypatch): + monkeypatch.setattr(actions, "execute", fake_execute(exc=actions.ChallengeFailed("cloudflare not passed: timeout"))) + data = post(env, screenshot=True).json() + assert data["status"] == "err" + assert data["code"] == "challenge_failed" + assert data["error"] == "cloudflare not passed: timeout" + assert data["actions"] == [] + assert data["cookies"] and data["screenshot"] + + +def test_request_failed(env): + FakeSession.error = RuntimeError("net::ERR_PROXY_CONNECTION_FAILED") + data = post(env).json() + assert data == { + "status": "err", "code": "request_failed", "error": "net::ERR_PROXY_CONNECTION_FAILED", + "actions": [], "duration_ms": data["duration_ms"], + } + + +@pytest.mark.parametrize("workers", ["busy", "none"]) +def test_no_slots(env, workers): + if workers == "busy": + env.heartbeat(free_slots=0) + else: + env.redis.flushall() + resp = post(env) + assert resp.status_code == 503 + assert resp.headers["retry-after"] == main.RETRY_AFTER + assert resp.json() == {"status": "err", "code": "no_slots", "error": "No slots available"} + assert env.messages == [] + + +def test_not_picked_up_in_time_is_503_and_cancelled(env, monkeypatch): + env.auto = False + monkeypatch.setattr(main, "V1_QUEUE_WAIT", 0) + resp = post(env) + assert resp.status_code == 503 + assert resp.json()["code"] == "no_slots" + + # воркер позже достаёт сообщение — отменённую задачу не выполняет + env.deliver_pending() + assert env.outcomes == [Outcome.SKIPPED] + assert FakeSession.instances == [] + + +def test_slot_released_after_request(env): + FakeSession.error = RuntimeError("boom") + for _ in range(3): + assert post(env).status_code == 200 + assert env.pool.used == 0 and env.pool.windows == {} + assert all(s.closed for s in FakeSession.instances) + + +def test_invalid_request_is_422_without_task(env): + resp = post(env, actions=[{"type": "nope"}]) + assert resp.status_code == 422 + assert env.messages == [] + + +def test_session_cookie_reused(env, tmp_path): + sid = str(uuid.uuid4()) + env.client.cookies.set("session_id", sid) + resp = post(env) + assert resp.cookies["session_id"] == sid + assert FakeSession.instances[0].kwargs["user_data_dir"] == str(tmp_path / sid) + + +@pytest.mark.parametrize("bad", ["../../etc", "not-a-uuid", "/tmp/x"]) +def test_bad_session_cookie_gets_new_session(env, tmp_path, bad): + env.client.cookies.set("session_id", bad) + resp = post(env) + sid = resp.cookies["session_id"] + assert sid != bad and uuid.UUID(sid) + assert [p.name for p in tmp_path.iterdir()] == [sid] + + +def test_user_agent_param_rejected(env): + resp = post(env, user_agent="First/1") + assert resp.status_code == 422 + assert "take user_agent from the response" in resp.text + assert env.messages == [] + + +def test_fingerprint_passed_to_browser_and_ua_untouched(env): + post(env) + kwargs = FakeSession.instances[0].kwargs + assert "useragent" not in kwargs + assert set(kwargs["additional_args"]) == {"screen", "device_scale_factor"} + + +def test_cookies_replace_profile_cookies_only_when_passed(env): + post(env) + post(env, cookies=[{"name": "a", "value": "1", "url": "https://example.com"}]) + post(env, cookies=[]) + assert [s.kwargs["clear_cookies"] for s in FakeSession.instances] == [False, True, True] + + +def test_proxy_passed_and_geo_resolved(env, monkeypatch): + seen = {} + + async def fake_geo(proxy): + seen["proxy"] = proxy + return {"locale": "de-DE", "timezone_id": "Europe/Berlin"} + + monkeypatch.setattr(windows, "resolve_geo", fake_geo) + post(env, proxy="http://u:p@host:3128") + kwargs = FakeSession.instances[0].kwargs + assert kwargs["proxy"] == {"server": "http://host:3128", "username": "u", "password": "p"} + assert seen["proxy"] == kwargs["proxy"] + assert (kwargs["locale"], kwargs["timezone_id"]) == ("de-DE", "Europe/Berlin") + + +def test_without_proxy_no_geo(env, monkeypatch): + async def fail_geo(proxy): + raise AssertionError("geo must not be resolved without proxy") + + monkeypatch.setattr(windows, "resolve_geo", fail_geo) + post(env) + assert FakeSession.instances[0].kwargs["proxy"] is None + + +def test_secrets_not_logged_nor_stored(env, monkeypatch, caplog): + async def no_geo(proxy): + return None + + monkeypatch.setattr(windows, "resolve_geo", no_geo) + caplog.set_level(logging.DEBUG) + post( + env, proxy="http://someuser:TopSecret1@host:3128", captcha="KEYSECRET2", html=True, + cookies=[{"name": "sid", "value": "CookieSecret3", "url": "https://example.com"}], + actions=[{"type": "fill", "selector": "#p", "value": "PassSecret4"}], + ) + log = caplog.text + assert "Incoming payload" in log and "Outgoing payload" in log + for secret in ("someuser", "TopSecret1", "KEYSECRET2", "CookieSecret3", "PassSecret4", "secret-cookie", "page"): + assert secret not in log + + # в Redis — только результат (там cookies и html страницы), но не секреты запроса + stored = " ".join(str(env.redis.hgetall(k)) for k in env.redis.keys("*task*")) + for secret in ("TopSecret1", "KEYSECRET2", "PassSecret4"): + assert secret not in stored + + +def test_health(env): + data = env.client.get("/health").json() + assert data["status"] == "ok" + assert (data["browser_slots"], data["free_slots"]) == (2, 2) + assert [w["worker_id"] for w in data["workers"]] == ["w1"] diff --git a/tests/unit/test_v2.py b/tests/unit/test_v2.py new file mode 100644 index 0000000..856c497 --- /dev/null +++ b/tests/unit/test_v2.py @@ -0,0 +1,180 @@ +"""/v2: задачи через очередь, доступ по cookie сессии, удержание окна и маршрутизация в него (#80–#83).""" +import uuid + +import pytest + +from tasks import store as task_store +from tasks.broker import COMMON_QUEUE, personal_queue +from worker.windows import Outcome + +from .conftest import FakeSession, WORKER_ID + + +def create(env, **body): + return env.client.post("/v2/tasks", json={"actions": [], **body}) + + +def get(env, task_id): + return env.client.get(f"/v2/tasks/{task_id}") + + +def session_key(sid): + return f"{task_store.PREFIX}:session:{sid}" + + +def test_create_and_get_result(env): + resp = create(env, url="https://example.com", html=True) + assert resp.status_code == 202 + body = resp.json() + assert body["status"] == "queued" and body["url"] == f"/v2/tasks/{body['task_id']}" + assert env.messages[0].queue_name == COMMON_QUEUE + + task = get(env, body["task_id"]).json() + assert task["status"] == "done" + assert task["result"]["status"] == "ok" + assert task["result"]["html"] == "page" + assert task["result"]["user_agent"] == "UA/engine" + + +def test_queued_until_worker_takes_it(env): + env.auto = False + task_id = create(env, url="https://example.com").json()["task_id"] + assert get(env, task_id).json() == {"task_id": task_id, "status": "queued"} + env.deliver_pending() + assert get(env, task_id).json()["status"] == "done" + + +@pytest.mark.parametrize("cookie", [None, "other", "garbage"]) +def test_task_visible_only_to_its_session(env, cookie): + task_id = create(env, url="https://example.com").json()["task_id"] + env.client.cookies.clear() + if cookie: + env.client.cookies.set("session_id", str(uuid.uuid4()) if cookie == "other" else cookie) + assert get(env, task_id).status_code == 404 + + +@pytest.mark.parametrize("task_id", [str(uuid.uuid4()), "../etc", "x"]) +def test_unknown_task_404(env, task_id): + create(env, url="https://example.com") + assert get(env, task_id).status_code == 404 + + +def test_result_ttl(env): + task_id = create(env, url="https://example.com").json()["task_id"] + ttl = env.redis.ttl(f"{task_store.PREFIX}:task:{task_id}") + assert task_store.RESULT_TTL - 5 < ttl <= task_store.RESULT_TTL + + +def test_url_required_without_window(env): + resp = create(env) + assert resp.status_code == 422 + assert "no open browser window" in resp.text + assert env.messages == [] + + +@pytest.mark.parametrize("value", [-1, 121]) +def test_idle_timeout_range(env, value): + assert create(env, url="https://example.com", browser_idle_timeout=value).status_code == 422 + + +def test_window_closed_by_default(env): + create(env, url="https://example.com") + sid = env.client.cookies["session_id"] + assert env.pool.windows == {} and env.pool.used == 0 + assert FakeSession.instances[0].closed + assert not env.redis.exists(session_key(sid)) + + +def test_window_held_and_reused(env): + create(env, url="https://example.com", browser_idle_timeout=5) + sid = env.client.cookies["session_id"] + engine = FakeSession.instances[0] + assert not engine.closed and env.pool.used == 1 + assert env.redis.hgetall(session_key(sid)) == {"worker_id": WORKER_ID, "idle_timeout": "5"} + assert 290 < env.redis.ttl(session_key(sid)) <= 300 + + # следующая задача сессии — в персональную очередь воркера с окном, без url — на текущей странице + task_id = create(env, browser_idle_timeout=2).json()["task_id"] + assert env.messages[-1].queue_name == personal_queue(WORKER_ID) + assert get(env, task_id).json()["result"]["status"] == "ok" + assert FakeSession.instances == [engine] + assert engine.calls == [("fetch", "https://example.com"), ("run", None)] + # простой продлён значением из последней задачи + assert 110 < env.redis.ttl(session_key(sid)) <= 120 + + # задача с browser_idle_timeout=0 закрывает окно после себя + create(env, url="https://example.com/next") + assert engine.closed and env.pool.used == 0 + assert not env.redis.exists(session_key(sid)) + + +def test_cookies_replace_in_open_window(env): + create(env, url="https://example.com", browser_idle_timeout=5) + create(env, browser_idle_timeout=5, cookies=[{"name": "a", "value": "1", "url": "https://example.com"}]) + create(env, browser_idle_timeout=5) + assert FakeSession.instances[0].context.ops == [ + "clear", ("add", [{"name": "a", "value": "1", "url": "https://example.com"}]), + ] + + +def test_proxy_change_restarts_window(env): + create(env, url="https://example.com", browser_idle_timeout=5) + create(env, url="https://example.com", browser_idle_timeout=5, proxy="http://host:3128") + first, second = FakeSession.instances + assert first.closed and not second.closed + assert second.kwargs["proxy"] == {"server": "http://host:3128"} + assert env.pool.used == 1 + + +def test_failure_in_window_closes_it(env): + create(env, url="https://example.com", browser_idle_timeout=5) + sid = env.client.cookies["session_id"] + FakeSession.error = RuntimeError("net::ERR_TIMED_OUT") + task_id = create(env, url="https://example.com", browser_idle_timeout=5).json()["task_id"] + assert get(env, task_id).json()["result"]["code"] == "request_failed" + assert FakeSession.instances[0].closed and env.pool.used == 0 + assert not env.redis.exists(session_key(sid)) + + +def test_idle_window_expires(env): + create(env, url="https://example.com", browser_idle_timeout=5) + sid = env.client.cookies["session_id"] + env.call(env.pool.expire_idle) + assert env.pool.used == 1 + + env.redis.delete(session_key(sid)) # TTL истёк + env.call(env.pool.expire_idle) + assert env.pool.used == 0 and FakeSession.instances[0].closed + + +def test_delete_session_closes_window(env): + create(env, url="https://example.com", browser_idle_timeout=5) + sid = env.client.cookies["session_id"] + assert env.client.delete("/v2/sessions/current").json() == {"closed": True} + assert env.messages[-1].actor_name == "close_window" + assert env.messages[-1].queue_name == personal_queue(WORKER_ID) + assert FakeSession.instances[0].closed and env.pool.used == 0 + assert not env.redis.exists(session_key(sid)) + assert env.client.delete("/v2/sessions/current").json() == {"closed": False} + + +def test_open_windows_take_slots(env): + for _ in range(2): + env.client.cookies.clear() + create(env, url="https://example.com", browser_idle_timeout=5) + assert env.pool.used == 2 + + env.client.cookies.clear() + create(env, url="https://example.com") + assert env.outcomes[-1] == Outcome.NO_SLOT + # /v1/solve без окна при занятых слотах — сразу 503, как раньше + env.client.cookies.clear() + assert env.client.post("/v1/solve", json={"url": "https://example.com", "actions": []}).status_code == 503 + + +def test_dead_worker_window_goes_to_common_queue(env): + create(env, url="https://example.com", browser_idle_timeout=5) + env.redis.delete(f"{task_store.PREFIX}:worker:{WORKER_ID}") + env.auto = False + create(env, url="https://example.com") + assert env.messages[-1].queue_name == COMMON_QUEUE