robotex/src/tasks/store.py

"""
Состояние задач, окон и воркеров в Redis (#78). Сроки жизни — TTL ключей Redis.

* `robotex:task:<id>` (hash) — статус задачи, сессия, время, результат. Пока задача не завершена,
  у ключа страховочный TTL `PENDING_TTL`; при завершении — `RESULT_TTL` (`RESULT_TTL_HOURS`).
* `robotex:session:<id>` (hash) — где открыто окно сессии: `worker_id` и `idle_timeout`.
  Пока задача выполняется, TTL — `RUNNING_TTL`; после неё — `browser_idle_timeout`. Ключ исчез —
  окно закрывается (worker/windows.py).
* `robotex:worker:<id>` — 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"])