spb/company-atlas

1"""`catlas schedule` — the long-running crawl process (spec §15–16, §62–63, §124–128).23Every `settings.scheduler_tick_s` the scheduler claims due sensors from Postgres (`for update skip locked`, so any number of4scheduler processes on any machines can share one database), runs them through `pipeline.run_sensor` under a global semaphore5(the per-domain governor lives in `fetch.Fetcher`), releases claims, writes a `crawl_runs` row and a heartbeat in `settings_kv`.6It also hosts the onboarding worker (`discovery.onboard_pending`) and every periodic task registered through7`services.periodic` (interval + cron via APScheduler in `settings.tz`). Claims older than `CLAIM_TTL` are considered abandoned.8"""9from __future__ import annotations1011import asyncio12import logging13import os14import signal15import socket16import time17from datetime import UTC, datetime18from typing import Any1920from companyatlas.config import settings21from companyatlas.db import execute, fetch_all, jsonb, transaction22from companyatlas.fetch import Fetcher23from companyatlas.ids import new_id24from companyatlas.services import periodic25from companyatlas.services.pipeline import RunOutcome, run_sensor2627log = logging.getLogger(__name__)2829CLAIM_TTL = "15 minutes"30IDLE_RUN_ROW_EVERY = 40            # write a crawl_runs row on idle ticks only every N ticks31ONBOARD_BATCH = 2032ONBOARD_EVERY_TICKS = 4333435def worker_name() -> str:36    return f"{socket.gethostname().split('.')[0]}:{os.getpid()}"373839async def claim_due_sensors(worker: str, limit: int) -> list[dict[str, Any]]:40    async with transaction() as conn:41        return await fetch_all(conn, f"""42            with due as (43                select id from sensors44                where status in ('active', 'failing', 'pending') and next_run_at <= now()45                  and (claimed_at is null or claimed_at < now() - interval '{CLAIM_TTL}')46                order by priority desc, next_run_at47                limit :limit for update skip locked)48            update sensors s set claimed_by = :worker, claimed_at = now() from due where s.id = due.id returning s.*49        """, limit=limit, worker=worker)505152async def release_claims(worker: str) -> None:53    async with transaction() as conn:54        await execute(conn, "update sensors set claimed_by = null, claimed_at = null where claimed_by = :w", w=worker)555657async def heartbeat(worker: str, payload: dict[str, Any]) -> None:58    async with transaction() as conn:59        for key in ("scheduler:heartbeat", f"scheduler:heartbeat:{worker}"):60            await execute(conn, """insert into settings_kv (key, value) values (:k, cast(:v as jsonb))61                                   on conflict (key) do update set value = excluded.value, updated_at = now()""", k=key, v=jsonb(payload))626364async def record_run(kind: str, worker: str, started_at: datetime, stats: dict[str, Any], error: str | None = None) -> None:65    async with transaction() as conn:66        await execute(conn, """insert into crawl_runs (id, kind, worker, started_at, finished_at, stats, error)67                               values (:id, :kind, :worker, :started, now(), cast(:stats as jsonb), :error)""",68                      id=new_id("crawl_run"), kind=kind, worker=worker, started=started_at, stats=jsonb(stats), error=error)697071class Scheduler:72    def __init__(self, *, concurrency: int | None = None, onboarding: bool = True, worker: str | None = None, claim_batch: int | None = None):73        self.concurrency = max(1, concurrency or settings.fetch_concurrency)74        self.onboarding = onboarding75        self.worker = worker or worker_name()76        self.claim_batch = claim_batch or settings.scheduler_claim_batch77        self.stop = asyncio.Event()78        self.inflight = 079        self.ticks = 080        self.fetcher = Fetcher()81        self._onboard_task: asyncio.Task[Any] | None = None82        self._aps: Any = None8384    # ------------------------------------------------------------------ one tick85    async def tick(self) -> dict[str, Any]:86        started = datetime.now(UTC)87        t0 = time.perf_counter()88        sensors = await claim_due_sensors(self.worker, self.claim_batch)89        stats = {"claimed": len(sensors), "ok": 0, "changed": 0, "meaningful": 0, "failed": 0, "not_modified": 0, "unchanged": 0, "skipped": 0, "redirected": 0}90        sem = asyncio.Semaphore(self.concurrency)9192        async def one(row: dict[str, Any]) -> None:93            async with sem:94                self.inflight += 195                try:96                    outcome: RunOutcome = await run_sensor(row, fetcher=self.fetcher, worker=self.worker)97                except Exception as exc:98                    log.exception("sensor run crashed", extra={"sensor_id": row.get("id")})99                    stats["failed"] += 1100                    async with transaction() as conn:101                        await execute(conn, """update sensors set claimed_by = null, claimed_at = null, last_error = :e,102                                               next_run_at = now() + interval '30 minutes' where id = :id""", id=row["id"], e=f"{exc.__class__.__name__}: {exc}"[:500])103                    return104                finally:105                    self.inflight -= 1106                if outcome.status in ("ok", "changed", "unchanged", "not_modified"):107                    stats["ok"] += 1108                if outcome.status == "changed":109                    stats["changed"] += 1110                    if outcome.kind in ("meaningful", "major", "critical"):111                        stats["meaningful"] += 1112                elif outcome.status in stats:113                    stats[outcome.status] += 1114115        if sensors:116            await asyncio.gather(*(one(s) for s in sensors))117        stats["tick_ms"] = int((time.perf_counter() - t0) * 1000)118        self.ticks += 1119        if sensors or self.ticks % IDLE_RUN_ROW_EVERY == 0:120            await record_run("scheduler_tick", self.worker, started, stats)121        due = await self._due_count()122        await heartbeat(self.worker, {"worker": self.worker, "at": datetime.now(UTC).isoformat(), "inflight": self.inflight, "due": due, "tick_ms": stats["tick_ms"],123                                      "claimed": stats["claimed"], "concurrency": self.concurrency, "tasks": periodic.snapshot()})124        log.info("tick", extra={"worker": self.worker, **stats, "due": due})125        return stats126127    async def _due_count(self) -> int:128        async with transaction() as conn:129            rows = await fetch_all(conn, "select count(*) as n from sensors where status in ('active','failing','pending') and next_run_at <= now()")130        return int(rows[0]["n"]) if rows else 0131132    # ------------------------------------------------------------------ onboarding133    async def _maybe_onboard(self) -> None:134        if not self.onboarding or (self._onboard_task and not self._onboard_task.done()):135            return136        from companyatlas.services.discovery import onboard_pending137138        async def go() -> None:139            try:140                stats = await onboard_pending(ONBOARD_BATCH, settings.onboarding_concurrency, fetcher=self.fetcher, worker=self.worker)141                if stats.get("claimed"):142                    await record_run("onboarding", self.worker, datetime.now(UTC), stats)143            except Exception:144                log.exception("onboarding batch failed")145146        self._onboard_task = asyncio.create_task(go())147148    # ------------------------------------------------------------------ periodic tasks149    def _start_periodic(self) -> None:150        loaded = periodic.load_task_modules()151        try:152            from apscheduler.schedulers.asyncio import AsyncIOScheduler153            from apscheduler.triggers.cron import CronTrigger154            from apscheduler.triggers.interval import IntervalTrigger155        except ImportError:  # pragma: no cover156            log.warning("apscheduler unavailable — periodic tasks disabled")157            return158        aps = AsyncIOScheduler(timezone=settings.tz)159        for task in periodic.tasks().values():160            if task.cron:161                trigger: Any = CronTrigger.from_crontab(task.cron, timezone=settings.tz)162            else:163                trigger = IntervalTrigger(seconds=float(task.every_s or 60), start_date=datetime.now(UTC).timestamp() + task.initial_delay_s and None)164            aps.add_job(task.run, trigger=trigger, id=task.name, name=task.name, max_instances=1, coalesce=True, misfire_grace_time=60)165        aps.start()166        self._aps = aps167        log.info("periodic tasks started", extra={"modules": loaded, "tasks": list(periodic.tasks())})168169    # ------------------------------------------------------------------ main loop170    async def run(self, *, once: bool = False) -> None:171        await self.fetcher.open()172        loop = asyncio.get_running_loop()173        for sig in (signal.SIGINT, signal.SIGTERM):174            try:175                loop.add_signal_handler(sig, self.stop.set)176            except (NotImplementedError, RuntimeError):  # pragma: no cover - non-unix177                pass178        try:179            from companyatlas.sdk.connector import sync_connectors_table180181            async with transaction() as conn:182                await sync_connectors_table(conn)183            if not once:184                self._start_periodic()185            log.info("scheduler started", extra={"worker": self.worker, "concurrency": self.concurrency, "tick_s": settings.scheduler_tick_s, "onboarding": self.onboarding})186            while not self.stop.is_set():187                try:188                    if self.onboarding and (self.ticks % ONBOARD_EVERY_TICKS == 0):189                        await self._maybe_onboard()190                    await self.tick()191                except Exception:192                    log.exception("tick failed")193                if once:194                    if self._onboard_task:195                        await self._onboard_task196                    break197                try:198                    await asyncio.wait_for(self.stop.wait(), settings.scheduler_tick_s)199                except TimeoutError:200                    pass201        finally:202            if self._aps is not None:203                self._aps.shutdown(wait=False)204            if self._onboard_task and not self._onboard_task.done():205                self._onboard_task.cancel()206            await release_claims(self.worker)207            await self.fetcher.close()208            log.info("scheduler stopped", extra={"worker": self.worker, "ticks": self.ticks})209210211async def run_scheduler(*, concurrency: int | None = None, onboarding: bool = True, once: bool = False, worker: str | None = None) -> None:212    await Scheduler(concurrency=concurrency, onboarding=onboarding, worker=worker).run(once=once)213214215__all__ = ["Scheduler", "claim_due_sensors", "heartbeat", "record_run", "release_claims", "run_scheduler", "worker_name"]216