protokeyandClaude Opus 5.5
502b6e1003
Два процесса на одном хосте получали один 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 <noreply@anthropic.com>
124 lines
5.4 KiB
Python
124 lines
5.4 KiB
Python
"""
|
||
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 signal
|
||
|
||
LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()
|
||
logging.basicConfig(
|
||
level=LOG_LEVEL,
|
||
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
|
||
)
|
||
# под dramatiq логирование уже настроено им самим (уровень INFO) и basicConfig ничего не делает — уровень явно
|
||
logging.getLogger().setLevel(LOG_LEVEL)
|
||
|
||
import dramatiq # noqa: E402
|
||
from dramatiq.asyncio import get_event_loop_thread # noqa: E402
|
||
from dramatiq.middleware import Middleware # noqa: E402
|
||
|
||
from db.engine import connect as connect_db # noqa: E402
|
||
from db.journal import Journal # 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__)
|
||
|
||
# сколько окон Chrome одновременно — и выполняющих задачу, и удерживаемых; ≈ числу ядер (docs/Deployment.md)
|
||
BROWSER_SLOTS = int(os.getenv("BROWSER_SLOTS", "4"))
|
||
# через сколько задача без свободного слота снова попадёт в очередь
|
||
NO_SLOT_DELAY_MS = 2000
|
||
HEARTBEAT_INTERVAL = 10
|
||
PROFILE_CLEANUP_INTERVAL = 3600
|
||
TASK_TIME_LIMIT_MS = task_store.RUNNING_TTL * 1000
|
||
|
||
_pool: WindowPool | None = None
|
||
_background: list[asyncio.Task] = []
|
||
|
||
|
||
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 _cleanup_loop(pool: WindowPool) -> None:
|
||
"""Удаление старых профилей: сразу при старте и затем раз в час."""
|
||
while True:
|
||
try:
|
||
await pool.cleanup_profiles()
|
||
except Exception:
|
||
log.warning("profile cleanup failed", exc_info=True)
|
||
await asyncio.sleep(PROFILE_CLEANUP_INTERVAL)
|
||
|
||
|
||
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()), 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))
|
||
return worker_id
|
||
|
||
|
||
async def _stop() -> None:
|
||
for task in _background:
|
||
task.cancel()
|
||
_background.clear()
|
||
await _pool.close_all()
|
||
await _pool.unregister()
|
||
if _pool.journal.db:
|
||
await _pool.journal.db.close()
|
||
|
||
|
||
class Lifecycle(Middleware):
|
||
"""Пул окон и heartbeat — в event loop AsyncIO-middleware: он поднимается раньше (middleware
|
||
добавлен в брокер раньше), а окна закрываются до его остановки."""
|
||
|
||
def before_worker_boot(self, broker, worker):
|
||
# персональная очередь объявляется до запуска потоков воркера: dramatiq подпишется и на неё
|
||
broker.declare_queue(personal_queue(_run(_start())))
|
||
|
||
def before_worker_shutdown(self, broker, worker):
|
||
_run(_stop())
|
||
|
||
|
||
broker.add_middleware(Lifecycle())
|
||
|
||
|
||
@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, meta: dict | None = None) -> None:
|
||
if _run(_pool.handle(task_id, session_id, payload, idle_timeout, meta)) == Outcome.NO_SLOT:
|
||
# в общую очередь: окно могло открыться на любом воркере, где есть слот
|
||
run_task.send_with_options(args=(task_id, session_id, payload, idle_timeout, meta), 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))
|