From 502b6e10030e3b73f45b19f02fa8ac5852f52a3d Mon Sep 17 00:00:00 2001 From: protokey Date: Wed, 14 Oct 2026 15:02:30 +0400 Subject: [PATCH] =?UTF-8?q?fix:=20=D0=B8=D0=BC=D1=8F=20=D0=B2=D0=BE=D1=80?= =?UTF-8?q?=D0=BA=D0=B5=D1=80=D0=B0=20=E2=80=94=20=D1=81=D0=B2=D0=BE=D0=B1?= =?UTF-8?q?=D0=BE=D0=B4=D0=BD=D0=BE=D0=B5=20worker-N=20=D0=B8=D0=B7=20Redi?= =?UTF-8?q?s,=20=D0=B0=20=D0=BD=D0=B5=20hostname=20(#110)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Два процесса на одном хосте получали один WORKER_ID и делили персональную очередь: задачу без url в открытое окно забирал процесс без этого окна. Теперь при старте воркер занимает наименьшее свободное имя worker-N (SET NX на heartbeat-ключ), heartbeat продлевает его только при своём instance; если имя занял другой процесс — воркер останавливается. Имена и персональные очереди переиспользуются. WORKER_ID из env убран, в /health — instance, host, pid. make dev при выходе гасит и api, и worker. Co-Authored-By: Claude Opus 5.5 --- Makefile | 7 ++- docs/Deployment.md | 6 ++- src/tasks/store.py | 54 +++++++++++++++++++--- src/worker/actors.py | 26 +++++++---- src/worker/windows.py | 37 +++++++++++++-- tests/unit/test_actors.py | 8 ++-- tests/unit/test_routing.py | 92 ++++++++++++++++++++++++++++++++++++++ tests/unit/test_store.py | 29 ++++++++++++ 8 files changed, 236 insertions(+), 23 deletions(-) create mode 100644 tests/unit/test_routing.py diff --git a/Makefile b/Makefile index e7f51b7..1c395dd 100644 --- a/Makefile +++ b/Makefile @@ -115,11 +115,14 @@ dev-deps: $(DEV_COMPOSE) up -d --wait rabbitmq redis postgres # Ctrl+C останавливает оба процесса; инфраструктура остаётся запущенной (make down — остановить) +# api и worker — в фоне: фоновые процессы sh не получают SIGINT от Ctrl+C, поэтому при выходе (Ctrl+C, закрытие +# терминала, kill) trap останавливает оба и ждёт их — иначе осиротевший worker остаётся работать (#110) dev: dev-deps cd src && export $(DEV_ENV) && \ { $(PYTHON) -m uvicorn main:app --reload --port $(DEV_API_PORT) & api=$$!; \ - trap 'kill $$api 2>/dev/null' EXIT INT TERM; \ - $(PYTHON) -m dramatiq worker.actors --processes 1 --threads $$(( $(DEV_SLOTS) * 2 )); } + $(PYTHON) -m dramatiq worker.actors --processes 1 --threads $$(( $(DEV_SLOTS) * 2 )) & worker=$$!; \ + trap 'trap - EXIT INT TERM HUP; kill $$api $$worker 2>/dev/null; wait' EXIT INT TERM HUP; \ + wait $$worker; } dev-key: dev-deps cd src && $(DEV_ENV) $(PYTHON) -m db.cli create "$(name)" diff --git a/docs/Deployment.md b/docs/Deployment.md index e7a8413..b3dda41 100644 --- a/docs/Deployment.md +++ b/docs/Deployment.md @@ -32,6 +32,11 @@ API доступно на `http://localhost:8000`. а задачи сессии с открытым окном api направляет в персональную очередь воркера, где это окно. Профили сессий лежат в volume воркера — воркеры на одном хосте делят его. +Имя воркера (`worker-1`, `worker-2`, …) не настраивается: при старте каждый процесс занимает в Redis наименьшее +свободное имя, поэтому два воркера на одном хосте не делят персональную очередь, а после перезапусков имена и +очереди переиспользуются. Если имя за время недоступности Redis занял другой процесс, воркер останавливается, +и Docker перезапускает его с новым именем. Где работает воркер — поля `host` и `pid` в `/health`. + ## Клиенты и ключи Без ключей API никого не пускает. Ключ показывается один раз при выпуске, в базе хранится только sha256. @@ -111,7 +116,6 @@ docker compose exec worker sh -c 'rm -rf /app/extra/user_data_dir/*' |---|---|---| | `BROWSER_SLOTS` | `4` | worker: сколько окон Chrome открыто одновременно — и выполняющих задачу, и удерживаемых. Задача без свободного слота ждёт в очереди | | `WORKER_THREADS` | `2 × BROWSER_SLOTS` | worker: потоков dramatiq | -| `WORKER_ID` | hostname | worker: имя воркера и его персональной очереди | | `RESULT_TTL_HOURS` | `24` | сколько хранится результат задачи | | `PROFILE_TTL_DAYS` | `7` | worker: через сколько дней без использования удаляется профиль браузера | | `V1_QUEUE_WAIT` | `30` | api: сколько `/v1/solve` ждёт, пока воркер возьмёт задачу; дальше `503` | diff --git a/src/tasks/store.py b/src/tasks/store.py index 929e493..bf5ef77 100644 --- a/src/tasks/store.py +++ b/src/tasks/store.py @@ -6,10 +6,13 @@ * `robotex:session:` (hash) — где открыто окно сессии: `worker_id`, `idle_timeout`, `client_id`. Пока задача выполняется, TTL — `RUNNING_TTL`; после неё — `browser_idle_timeout`. Ключ исчез — окно закрывается (worker/windows.py). -* `robotex:worker:` — heartbeat воркера с TTL `HEARTBEAT_TTL`: слоты и число окон. +* `robotex:worker:` — heartbeat воркера с TTL `HEARTBEAT_TTL`: слоты, число окон, `instance` (токен + процесса), `host`, `pid`. Ключ — он же аренда имени: имена `worker-1`, `worker-2`, … переиспользуются, + при старте воркер берёт наименьшее свободное (#110), поэтому персональных очередей не больше, чем воркеров. Прокси и ключ 2captcha сюда не попадают: они передаются только в сообщении очереди. """ +import itertools import json import os import time @@ -25,6 +28,7 @@ PENDING_TTL = 6 * 3600 # пока задача выполняется, ключ сессии держит маршрут к воркеру; дольше задача не живёт (TimeLimit) RUNNING_TTL = 3 * 3600 HEARTBEAT_TTL = 30 +WORKER_NAME_PREFIX = "worker-" # сколько раз задачу можно начать заново после падения воркера, прежде чем считать её проваленной MAX_ATTEMPTS = 2 @@ -186,11 +190,51 @@ class TaskStore: # --- воркеры --- - 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 acquire_worker_name(self, info: dict) -> str: + """Занять наименьшее свободное имя worker-N. info — heartbeat с instance процесса. + SET NX атомарен: два одновременно стартующих воркера одно имя не возьмут.""" + for n in itertools.count(1): + name = f"{WORKER_NAME_PREFIX}{n}" + if await self.redis.set(_worker_key(name), json.dumps(info), nx=True, ex=HEARTBEAT_TTL): + return name - async def remove_worker(self, worker_id: str) -> None: - await self.redis.delete(_worker_key(worker_id)) + async def heartbeat(self, worker_id: str, info: dict) -> bool: + """Продлить heartbeat, только если имя ещё за этим процессом (info["instance"]). + False — имя занял другой процесс: этот воркер его потерял.""" + key = _worker_key(worker_id) + async with self.redis.pipeline(transaction=True) as pipe: + while True: + try: + await pipe.watch(key) + current = await pipe.get(key) + if current is not None and json.loads(current).get("instance") != info.get("instance"): + await pipe.unwatch() + return False + pipe.multi() + pipe.set(key, json.dumps(info), ex=HEARTBEAT_TTL) + await pipe.execute() + return True + except WatchError: + continue + + async def remove_worker(self, worker_id: str, instance: str | None = None) -> None: + """Освободить имя. С instance — только если оно ещё за этим процессом.""" + key = _worker_key(worker_id) + if instance is None: + await self.redis.delete(key) + return + async with self.redis.pipeline(transaction=True) as pipe: + try: + await pipe.watch(key) + current = await pipe.get(key) + if current is None or json.loads(current).get("instance") != instance: + await pipe.unwatch() + return + pipe.multi() + pipe.delete(key) + await pipe.execute() + except WatchError: + pass async def worker_alive(self, worker_id: str | None) -> bool: return bool(worker_id) and bool(await self.redis.exists(_worker_key(worker_id))) diff --git a/src/worker/actors.py b/src/worker/actors.py index 759c4b8..9ed606b 100644 --- a/src/worker/actors.py +++ b/src/worker/actors.py @@ -8,7 +8,7 @@ dramatiq, а не из event loop. import asyncio import logging import os -import socket +import signal LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper() logging.basicConfig( @@ -30,7 +30,6 @@ 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")) # через сколько задача без свободного слота снова попадёт в очередь @@ -68,14 +67,23 @@ async def _cleanup_loop(pool: WindowPool) -> None: await asyncio.sleep(PROFILE_CLEANUP_INTERVAL) -async def _start() -> None: +def _stop_self() -> None: + """Имя воркера занял другой процесс: штатная остановка, как по Ctrl+C. dramatiq при выходе дочернего + процесса останавливается целиком, Docker (restart: unless-stopped) запустит воркер заново с новым именем.""" + os.kill(os.getpid(), signal.SIGTERM) + + +async def _start() -> str: + """Пул окон и имя воркера. Имя (worker-N) — не из env, а свободное в Redis (#110): два процесса на одном + хосте не делят персональную очередь, а имена и очереди переиспользуются после перезапусков.""" global _pool db = connect_db() await db.init() - _pool = WindowPool(task_store.TaskStore(task_store.connect()), WORKER_ID, BROWSER_SLOTS, Journal(db)) - await _pool.report() + _pool = WindowPool(task_store.TaskStore(task_store.connect()), None, BROWSER_SLOTS, Journal(db), on_lost=_stop_self) + worker_id = await _pool.register() _background.extend([asyncio.create_task(_heartbeat_loop(_pool)), asyncio.create_task(_cleanup_loop(_pool))]) - log.info("worker %s: %d browser slots, queues %s, %s", WORKER_ID, BROWSER_SLOTS, COMMON_QUEUE, personal_queue(WORKER_ID)) + log.info("worker %s: %d browser slots, queues %s, %s", worker_id, BROWSER_SLOTS, COMMON_QUEUE, personal_queue(worker_id)) + return worker_id async def _stop() -> None: @@ -83,7 +91,7 @@ async def _stop() -> None: task.cancel() _background.clear() await _pool.close_all() - await _pool.store.remove_worker(WORKER_ID) + await _pool.unregister() if _pool.journal.db: await _pool.journal.db.close() @@ -93,14 +101,14 @@ class Lifecycle(Middleware): добавлен в брокер раньше), а окна закрываются до его остановки.""" def before_worker_boot(self, broker, worker): - _run(_start()) + # персональная очередь объявляется до запуска потоков воркера: dramatiq подпишется и на неё + broker.declare_queue(personal_queue(_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) diff --git a/src/worker/windows.py b/src/worker/windows.py index ab75b11..b86921d 100644 --- a/src/worker/windows.py +++ b/src/worker/windows.py @@ -18,11 +18,13 @@ import base64 import logging import os import shutil +import socket import time +import uuid from contextlib import asynccontextmanager, suppress from dataclasses import dataclass, field from pathlib import Path -from typing import Any +from typing import Any, Callable from urllib.parse import urlsplit from patchright.async_api import Page, ProxySettings @@ -162,10 +164,18 @@ def build_result(payload: TaskRequest, result: dict) -> dict: class WindowPool: - def __init__(self, store: TaskStore, worker_id: str, slots: int, journal: Journal | None = None): + def __init__( + self, store: TaskStore, worker_id: str | None, slots: int, journal: Journal | None = None, + on_lost: Callable[[], None] | None = None, + ): + """worker_id None — имя берётся в Redis при register() (#110). + on_lost — вызывается, если имя занял другой процесс: такому воркеру надо остановиться.""" self.store = store self.journal = journal or Journal(None) self.worker_id = worker_id + self.instance = uuid.uuid4().hex + self.on_lost = on_lost + self.lost = False self.slots = slots self.used = 0 self.windows: dict[str, Window] = {} @@ -178,8 +188,19 @@ class WindowPool: "free_slots": self.slots - self.used, "windows": len(self.windows), "running": len(self.running), + "instance": self.instance, + "host": socket.gethostname(), + "pid": os.getpid(), } + async def register(self) -> str: + """Занять свободное имя worker-N: оно же — персональная очередь воркера.""" + self.worker_id = await self.store.acquire_worker_name(self.info()) + return self.worker_id + + async def unregister(self) -> None: + await self.store.remove_worker(self.worker_id, self.instance) + @asynccontextmanager async def _session_lock(self, session_id: str): entry = self._locks.setdefault(session_id, _SessionLock()) @@ -194,10 +215,20 @@ class WindowPool: async def report(self) -> None: """heartbeat со слотами сразу после их изменения: по нему api решает, отвечать ли 503 на /v1/solve.""" + if self.lost: + return try: - await self.store.heartbeat(self.worker_id, self.info()) + owned = await self.store.heartbeat(self.worker_id, self.info()) except Exception: log.warning("heartbeat failed", exc_info=True) + return + if not owned: + # Redis был недоступен дольше HEARTBEAT_TTL, и имя занял другой процесс: задачи его очереди + # теперь не наши — останавливаемся, после перезапуска возьмём свободное имя + self.lost = True + log.error("worker name %s is taken by another process, stopping", self.worker_id) + if self.on_lost: + self.on_lost() def _busy(self, session_id: str) -> bool: entry = self._locks.get(session_id) diff --git a/tests/unit/test_actors.py b/tests/unit/test_actors.py index 70fdc29..dc8ed35 100644 --- a/tests/unit/test_actors.py +++ b/tests/unit/test_actors.py @@ -51,15 +51,17 @@ def submit(actors, server, queue, idle=0): @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)) + # имя воркера — первое свободное в Redis (#110), персональная очередь объявлена до запуска потоков + assert actors._pool.worker_id == "worker-1" + redis = submit(actors, server, COMMON_QUEUE if queue == "common" else personal_queue("worker-1")) 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}") + assert redis.exists("robotex:worker:worker-1") 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 redis.hgetall(f"robotex:session:{SID}")["worker_id"] == "worker-1" assert not FakeSession.instances[0].closed diff --git a/tests/unit/test_routing.py b/tests/unit/test_routing.py new file mode 100644 index 0000000..587aedd --- /dev/null +++ b/tests/unit/test_routing.py @@ -0,0 +1,92 @@ +"""Два процесса воркера на одном брокере и Redis (#110): у каждого своё имя и своя персональная очередь, +задача без url в открытое окно всегда попадает к воркеру, который держит это окно.""" +import itertools + +import pytest + +import main +from db.journal import Journal +from tasks.broker import COMMON_QUEUE, RUN_TASK, personal_queue +from worker import windows + + +@pytest.fixture +def two_workers(env, monkeypatch): + """Два пула вместо пула харнесса; общую очередь разбирают по очереди, как два consumer'а RabbitMQ.""" + env.redis.delete("robotex:worker:w1") + pools = [windows.WindowPool(env.store, None, slots=2, journal=Journal(env.db)) for _ in range(2)] + for pool in pools: + env.call(pool.register) + turn = itertools.cycle(pools) + delivered = [] + + async def dispatch(message): + if message.queue_name == COMMON_QUEUE: + pool = next(turn) + else: + pool = next(p for p in pools if personal_queue(p.worker_id) == message.queue_name) + delivered.append((message.args[1], pool.worker_id)) + if message.actor_name == RUN_TASK: + await pool.handle(*message.args) + else: + await pool.close(*message.args) + await pool.report() + + monkeypatch.setattr(main, "dispatch", dispatch) + yield pools, delivered + for pool in pools: + env.call(pool.close_all) + + +def run(env, session_id=None, **body): + env.client.cookies.clear() + if session_id: + env.client.cookies.set("session_id", session_id) + resp = env.client.post("/v2/tasks", json={"actions": [], "browser_idle_timeout": 5, **body}) + assert resp.status_code == 202, resp.text + task = env.client.get(f"/v2/tasks/{resp.json()['task_id']}").json() + return resp.cookies.get("session_id") or session_id, task + + +def test_workers_get_distinct_names(env, two_workers): + pools, _ = two_workers + assert [p.worker_id for p in pools] == ["worker-1", "worker-2"] + workers = env.client.get("/health").json()["workers"] + assert [w["worker_id"] for w in workers] == ["worker-1", "worker-2"] + assert len({w["instance"] for w in workers}) == 2 + + +def test_task_without_url_goes_to_window_owner(env, two_workers): + pools, delivered = two_workers + # окна двух сессий открываются на разных воркерах + sid_a, task = run(env, url="https://example.com/a") + assert task["result"]["status"] == "ok" + sid_b, task = run(env, url="https://example.com/b") + assert task["result"]["status"] == "ok" + owner = {sid: wid for sid, wid in delivered} + assert sorted(owner.values()) == ["worker-1", "worker-2"] + + # следующие задачи без url — всегда в окно своей сессии + for _ in range(3): + for sid in (sid_a, sid_b): + _, task = run(env, sid) + assert task["status"] == "done" and task["result"]["status"] == "ok", task + assert delivered[-1] == (sid, owner[sid]) + assert all(len(p.windows) == 1 for p in pools) + + +def test_worker_stops_when_name_is_taken(env, two_workers): + """Имя истекло (Redis был недоступен) и досталось другому процессу — старый владелец останавливается.""" + pools, _ = two_workers + stopped = [] + pools[0].on_lost = lambda: stopped.append(True) + env.redis.delete("robotex:worker:worker-1") + newcomer = windows.WindowPool(env.store, None, slots=1) + assert env.call(newcomer.register) == "worker-1" + + env.call(pools[0].report) + env.call(pools[0].report) + assert stopped == [True] and pools[0].lost + # при остановке старый владелец не снимает чужой heartbeat + env.call(pools[0].unregister) + assert env.client.get("/health").json()["workers"][0]["instance"] == newcomer.instance diff --git a/tests/unit/test_store.py b/tests/unit/test_store.py index 6d35d64..c939eb9 100644 --- a/tests/unit/test_store.py +++ b/tests/unit/test_store.py @@ -72,3 +72,32 @@ async def test_live_window_requires_live_worker(store): await store.heartbeat("w1", {"browser_slots": 1, "free_slots": 0}) assert await store.live_window("s") == {"worker_id": "w1", "idle_timeout": "5", "client_id": "7"} assert await store.workers() == [{"worker_id": "w1", "browser_slots": 1, "free_slots": 0}] + + +async def test_worker_names_are_unique_and_reused(store): + """Имена worker-N: занятое не выдаётся, освободившееся — переиспользуется (#110).""" + assert await store.acquire_worker_name({"instance": "a"}) == "worker-1" + assert await store.acquire_worker_name({"instance": "b"}) == "worker-2" + await store.remove_worker("worker-1", "a") + assert await store.acquire_worker_name({"instance": "c"}) == "worker-1" + assert await store.redis.ttl("robotex:worker:worker-1") == task_store.HEARTBEAT_TTL + + +async def test_heartbeat_of_taken_name_is_refused(store): + """Имя истекло и его занял другой процесс: старый владелец его не продлит и не удалит.""" + assert await store.acquire_worker_name({"instance": "old"}) == "worker-1" + assert await store.heartbeat("worker-1", {"instance": "old", "windows": 1}) + await store.redis.delete("robotex:worker:worker-1") # TTL истёк, пока Redis был недоступен + assert await store.acquire_worker_name({"instance": "new"}) == "worker-1" + + assert not await store.heartbeat("worker-1", {"instance": "old"}) + await store.remove_worker("worker-1", "old") + assert (await store.workers())[0]["instance"] == "new" + + +async def test_expired_own_name_is_renewed(store): + """Имя истекло, но никто его не занял — heartbeat просто продлевает его.""" + assert await store.acquire_worker_name({"instance": "a"}) == "worker-1" + await store.redis.delete("robotex:worker:worker-1") + assert await store.heartbeat("worker-1", {"instance": "a"}) + assert await store.worker_alive("worker-1")