robotex/src/worker/windows.py

"""
Слоты браузера и окна сессий воркера (#81, #82).

Слот — один открытый Chrome. `BROWSER_SLOTS` ограничивает число окон воркера: и тех, где идёт задача,
и удерживаемых между задачами. Окно сессии остаётся открытым после задачи, если в ней задан
`browser_idle_timeout > 0`: следующая задача с той же cookie `session_id` попадает в это окно
(маршрутизация — в api, #83) и продолжает работу без перезапуска Chrome. Время простоя считает
Redis: TTL ключа `session:<id>`, каждая задача в окно ставит его заново. Ключ исчез — окно
закрывается фоновой проверкой (`expire_idle`).

Задачи одной сессии выполняются строго по очереди: профиль Chrome нельзя открыть дважды.

Всё здесь работает в одном event loop (поток AsyncIO-middleware dramatiq), поэтому счётчик слотов
и словарь окон не требуют блокировок между потоками.
"""
import asyncio
import base64
import logging
import 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, Callable
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_data, safe_request_dump, safe_response_dump
from api.schemas import TaskRequest
from db.journal import Journal
from engine import AsyncStealthySession
from tasks.sessions import parse_session_id
from tasks.store import RUNNING_TTL, Status, TaskStore

log = logging.getLogger(__name__)

USER_DATA_ROOT = (Path(__file__).parent / "../../extra/user_data_dir").resolve()
# профиль, которым не пользовались столько дней, удаляется (#20)
PROFILE_TTL_DAYS = float(os.getenv("PROFILE_TTL_DAYS", "7"))
# время последнего использования профиля — mtime этого файла
LAST_USED_FILE = ".robotex_last_used"


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
	client_id: int | None = None
	# строка окна в журнале (#17): открыто — закрыто
	journal_id: int | None = None


@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 touch_profile(user_data_dir: Path) -> None:
	(user_data_dir / LAST_USED_FILE).touch()


def profile_last_used(user_data_dir: Path) -> float:
	try:
		return (user_data_dir / LAST_USED_FILE).stat().st_mtime
	except OSError:
		# профиль из версии без отметки — по времени изменения каталога
		return user_data_dir.stat().st_mtime


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"),
			# адрес страницы после действий: SPA меняет его без навигации (номер заявки и т.п.)
			"url": page.url,
		}

	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"],
		"url": result.get("url"),
		"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 | 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] = {}
		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),
			"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())
		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."""
		if self.lost:
			return
		try:
			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)
		return bool(entry and entry.lock.locked())

	async def handle(
		self, task_id: str, session_id: str, payload: dict, idle_timeout: int, meta: dict | None = None,
	) -> str:
		"""Выполнить задачу. Результат — в Redis; возврат — что делать с сообщением очереди.
		meta: client_id и api (v1/v2) — для журнала."""
		meta = meta or {}
		client_id = meta.get("client_id")
		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, client_id=client_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)
			task = TaskRequest.model_validate(payload)
			result = None
			queued = await self.store.get(task_id) or {}
			await self.journal.request_started(
				task_id, client_id, session_id, meta.get("api", "v2"),
				float(queued["created_at"]) if queued.get("created_at") else None, safe_request_data(task),
			)
			try:
				await self.store.set_session(session_id, self.worker_id, idle_timeout, RUNNING_TTL, client_id)
				result = await self._execute(window, task, task_id)
				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)
				await self.journal.request_finished(task_id, window.journal_id, result)
				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, client_id)
				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, task_id: str) -> dict:
		started = time.monotonic()
		session_id = window.session_id
		log.info("Incoming request [%s]: %s", session_id, payload.url or "<current page>")
		log.debug("Incoming payload [%s]:\n%s", session_id, safe_request_dump(payload))

		page_action = make_page_action(payload)
		nav_kwargs = dict(
			timeout=payload.timeout * 1000 if payload.timeout else 60_000,
			disable_resources=payload.disable_resources,
			blocked_domains=set(payload.blocked_domains) if payload.blocked_domains else None,
		)
		try:
			if window.engine and window.proxy != payload.proxy:
				log.info("proxy changed, restarting browser [%s]", session_id)
				await self._close_engine(window)

			if window.engine is None:
				if not payload.url:
					raise NoWindow("url is required: the session has no open browser window")
				window.engine = await self._open(session_id, payload, nav_kwargs)
				window.proxy = payload.proxy
				window.journal_id = await self.journal.window_opened(window.client_id, session_id, task_id, self.worker_id)
			else:
				await self.journal.window_used(window.journal_id)
				if 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)
		touch_profile(user_data_dir)
		# отпечаток задаётся профилю один раз (см. api/profile.py)
		profile_kwargs = profile.browser_kwargs(profile.load_or_create(user_data_dir))
		# без менеджера паролей, перевода и автозаполнения (#111)
		profile.apply_chrome_prefs(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
		journal_id, window.journal_id = window.journal_id, None
		await self.journal.window_closed(journal_id)
		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:
		"""Закрыть окно и освободить его слот."""
		if window.engine is not None:
			with suppress(OSError):
				touch_profile(session_dir(window.session_id))
		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)

	async def cleanup_profiles(self, ttl_days: float = PROFILE_TTL_DAYS) -> list[str]:
		"""Удалить профили, которыми не пользовались ttl_days дней (#20). Профиль с открытым окном —
		на этом воркере или на другом (ключ session:<id> в Redis) — не трогается."""
		if not USER_DATA_ROOT.is_dir():
			return []
		deadline = time.time() - ttl_days * 86400
		removed = []
		for user_data_dir in USER_DATA_ROOT.iterdir():
			session_id = parse_session_id(user_data_dir.name)
			if not session_id or not user_data_dir.is_dir() or profile_last_used(user_data_dir) > deadline:
				continue
			if session_id in self.windows or self._busy(session_id) or await self.store.get_session(session_id):
				continue
			await asyncio.to_thread(shutil.rmtree, user_data_dir, True)
			removed.append(session_id)
		if removed:
			log.info("removed %d browser profiles unused for %s days", len(removed), ttl_days)
		return removed