iaai-parser/iaai_scraper/worker/self_heal.py

import json
import logging
import os
import random
import signal
import time

from redis import Redis


logger = logging.getLogger("iaai_scraper.worker.self_heal")

GLOBAL_PROGRESS_TS_KEY = "iaai:state:last_progress_ts"
SELF_HEAL_RESTART_LOCK_KEY = "iaai:state:self_heal_restart_in_progress"
SYNC_LISTING_LOCK_KEY = "iaai:locks:sync_listing"
SIGKILL_FALLBACK = getattr(signal, "SIGKILL", signal.SIGTERM)


def _env_bool(name: str, default: bool) -> bool:
	value = os.getenv(name)
	if value is None:
		return default
	return value.strip().lower() in {"1", "true", "yes", "on"}


def _env_int(name: str, default: int) -> int:
	value = os.getenv(name)
	if value is None:
		return default
	try:
		return int(value.strip())
	except Exception:
		return default


def _get_redis() -> Redis:
	url = os.getenv("IAAI_REDIS_URL", "redis://redis:6379/0")
	return Redis.from_url(
		url,
		decode_responses=True,
		socket_connect_timeout=5.0,
		socket_timeout=10.0,
		health_check_interval=30,
		retry_on_timeout=True,
	)


def _safe_int(value: str | None, default: int = 0) -> int:
	if value is None:
		return default
	try:
		return int(str(value).strip())
	except Exception:
		return default


def _read_last_progress_ts(redis_client: Redis) -> int | None:
	raw = redis_client.get(GLOBAL_PROGRESS_TS_KEY)
	if raw:
		ts = _safe_int(raw)
		if ts > 0:
			return ts

	# Fallback: если глобальный ключ не найден, берём max(ts) из task_progress:*.
	# Это дороже, но выполняется только при отсутствии основного маркера.
	max_ts = 0
	for key in redis_client.scan_iter(match="iaai:state:task_progress:*"):
		try:
			payload = redis_client.get(key)
			if not payload:
				continue
			data = json.loads(payload)
			ts = _safe_int(data.get("ts"), 0)
			if ts > max_ts:
				max_ts = ts
		except Exception:
			continue
	return max_ts or None


def _has_inflight_work(redis_client: Redis, queue_name: str) -> tuple[bool, dict[str, int]]:
	"""Есть ли признаки активной/зависшей работы, даже если очередь пуста."""
	queue_len = _safe_int(redis_client.llen(queue_name), 0)
	has_lock = 1 if redis_client.get(SYNC_LISTING_LOCK_KEY) else 0
	has_task_progress = 0
	for _ in redis_client.scan_iter(match="iaai:state:task_progress:*"):
		has_task_progress = 1
		break
	flags = {
		"queue_len": queue_len,
		"has_lock": has_lock,
		"has_task_progress": has_task_progress,
	}
	return (queue_len > 0 or has_lock == 1 or has_task_progress == 1), flags


def _kill_worker_process() -> None:
	pid_file = "/tmp/celery-worker.pid"
	pid: int | None = None
	try:
		with open(pid_file, "r", encoding="utf-8") as f:
			pid = int(f.read().strip())
	except Exception:
		pid = None

	if not pid:
		logger.error("Self-heal: failed to read worker pid from %s", pid_file)
		return

	logger.error("Self-heal: terminating stuck worker process pid=%s", pid)
	try:
		os.kill(pid, signal.SIGTERM)
	except Exception:
		logger.exception("Self-heal: failed to send SIGTERM to pid=%s", pid)
		return

	time.sleep(20)
	try:
		# Если процесс ещё жив — принудительно убиваем.
		os.kill(pid, 0)
		logger.error("Self-heal: worker pid=%s did not stop after SIGTERM; sending SIGKILL", pid)
		os.kill(pid, SIGKILL_FALLBACK)
	except ProcessLookupError:
		pass
	except Exception:
		logger.exception("Self-heal: failed to send SIGKILL to pid=%s", pid)


def main() -> None:
	if not _env_bool("IAAI_SELF_HEAL_ENABLED", True):
		logger.info("Self-heal watchdog disabled via IAAI_SELF_HEAL_ENABLED")
		return

	queue_name = os.getenv("IAAI_CELERY_QUEUE", "scraping")
	check_interval = max(5, _env_int("IAAI_SELF_HEAL_CHECK_INTERVAL_SECONDS", 30))
	stall_seconds = max(180, _env_int("IAAI_SELF_HEAL_STALL_SECONDS", 720))
	startup_grace = max(30, _env_int("IAAI_SELF_HEAL_STARTUP_GRACE_SECONDS", 300))
	restart_cooldown = max(60, _env_int("IAAI_SELF_HEAL_RESTART_COOLDOWN_SECONDS", 300))

	logger.warning(
		"Self-heal watchdog enabled: queue=%s check_interval=%ss stall=%ss startup_grace=%ss cooldown=%ss",
		queue_name,
		check_interval,
		stall_seconds,
		startup_grace,
		restart_cooldown,
	)

	started_at = time.time()
	redis_client: Redis | None = None

	while True:
		try:
			if redis_client is None:
				redis_client = _get_redis()
				redis_client.ping()

			has_inflight, inflight = _has_inflight_work(redis_client, queue_name)
			if not has_inflight:
				time.sleep(check_interval)
				continue

			last_progress_ts = _read_last_progress_ts(redis_client)
			now_ts = int(time.time())
			age = None if last_progress_ts is None else max(0, now_ts - int(last_progress_ts))

			if age is None:
				if now_ts - int(started_at) < startup_grace:
					time.sleep(check_interval)
					continue
				logger.warning(
					"Self-heal: inflight=%s but no progress timestamp found after startup grace",
					inflight,
				)
				age = stall_seconds + 1

			if age <= stall_seconds:
				time.sleep(check_interval)
				continue

			# Глобальный anti-storm lock: чтобы много воркеров не рестартились одновременно.
			acquired = bool(
				redis_client.set(
					SELF_HEAL_RESTART_LOCK_KEY,
					str(now_ts),
					nx=True,
					ex=restart_cooldown,
				)
			)
			if not acquired:
				time.sleep(check_interval)
				continue

			logger.error(
				"Self-heal: detected global stall (inflight=%s, progress_age=%ss > %ss). Restarting worker process...",
				inflight,
				age,
				stall_seconds,
			)
			# Небольшой джиттер, чтобы при одинаковом событии у разных контейнеров
			# перезапуск был не строго одновременно.
			time.sleep(random.uniform(0.3, 2.0))
			_kill_worker_process()
			# После kill pid1 контейнер будет перезапущен Docker restart-policy.
			# На случай неуспеха не молотим цикл.
			time.sleep(check_interval)

		except Exception:
			logger.exception("Self-heal watchdog iteration failed")
			redis_client = None
			time.sleep(check_interval)


if __name__ == "__main__":
	logging.basicConfig(
		level=os.getenv("IAAI_LOG_LEVEL", "INFO"),
		format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
	)
	main()