306 lines
11 KiB
Python
306 lines
11 KiB
Python
"""Task Scheduler for running collection jobs."""
|
|
|
|
import asyncio
|
|
import logging
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import Any, Dict, Optional
|
|
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from apscheduler.triggers.interval import IntervalTrigger
|
|
from sqlalchemy import select
|
|
|
|
from app.db.session import async_session_factory
|
|
from app.core.time import to_iso8601_utc
|
|
from app.models.datasource import DataSource
|
|
from app.models.task import CollectionTask
|
|
from app.services.collectors.registry import collector_registry
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
scheduler = AsyncIOScheduler()
|
|
RUNNING_TASK_GUARD_TIMEOUT_MINUTES = 90
|
|
RUNNING_COLLECTOR_TASKS: dict[str, asyncio.Task[Any]] = {}
|
|
|
|
|
|
def _collector_task_name(collector_name: str) -> str:
|
|
return f"collector:{collector_name}"
|
|
|
|
|
|
def get_running_collector_task(collector_name: str) -> asyncio.Task[Any] | None:
|
|
task = RUNNING_COLLECTOR_TASKS.get(collector_name)
|
|
if task is not None and not task.done():
|
|
return task
|
|
|
|
if task is not None and task.done():
|
|
RUNNING_COLLECTOR_TASKS.pop(collector_name, None)
|
|
|
|
target_name = _collector_task_name(collector_name)
|
|
for candidate in asyncio.all_tasks():
|
|
if candidate.done():
|
|
continue
|
|
if candidate.get_name() == target_name:
|
|
RUNNING_COLLECTOR_TASKS[collector_name] = candidate
|
|
return candidate
|
|
|
|
return None
|
|
|
|
|
|
async def _update_next_run_at(datasource: DataSource, session) -> None:
|
|
job = scheduler.get_job(datasource.source)
|
|
datasource.next_run_at = job.next_run_time if job else None
|
|
await session.commit()
|
|
|
|
|
|
async def _apply_datasource_schedule(datasource: DataSource, session) -> None:
|
|
collector = collector_registry.get(datasource.source)
|
|
if not collector:
|
|
logger.warning("Collector not found for datasource %s", datasource.source)
|
|
return
|
|
|
|
collector_registry.set_active(datasource.source, datasource.is_active)
|
|
|
|
existing_job = scheduler.get_job(datasource.source)
|
|
if existing_job:
|
|
scheduler.remove_job(datasource.source)
|
|
|
|
if datasource.is_active:
|
|
scheduler.add_job(
|
|
run_collector_task,
|
|
trigger=IntervalTrigger(minutes=max(1, datasource.frequency_minutes)),
|
|
id=datasource.source,
|
|
name=datasource.name,
|
|
replace_existing=True,
|
|
kwargs={"collector_name": datasource.source},
|
|
)
|
|
logger.info(
|
|
"Scheduled collector: %s (every %sm)",
|
|
datasource.source,
|
|
datasource.frequency_minutes,
|
|
)
|
|
else:
|
|
logger.info("Collector disabled: %s", datasource.source)
|
|
|
|
await _update_next_run_at(datasource, session)
|
|
|
|
|
|
async def run_collector_task(collector_name: str):
|
|
"""Run a single collector task."""
|
|
collector = collector_registry.get(collector_name)
|
|
if not collector:
|
|
logger.error("Collector not found: %s", collector_name)
|
|
return
|
|
|
|
async with async_session_factory() as db:
|
|
result = await db.execute(select(DataSource).where(DataSource.source == collector_name))
|
|
datasource = result.scalar_one_or_none()
|
|
if not datasource:
|
|
logger.error("Datasource not found for collector: %s", collector_name)
|
|
return
|
|
|
|
if not datasource.is_active:
|
|
logger.info("Skipping disabled collector: %s", collector_name)
|
|
return
|
|
|
|
running_result = await db.execute(
|
|
select(CollectionTask)
|
|
.where(
|
|
CollectionTask.datasource_id == datasource.id,
|
|
CollectionTask.status == "running",
|
|
)
|
|
.order_by(CollectionTask.started_at.desc(), CollectionTask.id.desc())
|
|
.limit(1)
|
|
)
|
|
existing_running = running_result.scalar_one_or_none()
|
|
if existing_running is not None:
|
|
now = datetime.now(UTC)
|
|
started_at = existing_running.started_at
|
|
if started_at is not None and started_at.tzinfo is None:
|
|
started_at = started_at.replace(tzinfo=UTC)
|
|
|
|
is_stale = (
|
|
started_at is not None
|
|
and (now - started_at) > timedelta(minutes=RUNNING_TASK_GUARD_TIMEOUT_MINUTES)
|
|
)
|
|
if not is_stale:
|
|
logger.warning(
|
|
"Skipping collector %s trigger because task %s is already running",
|
|
collector_name,
|
|
existing_running.id,
|
|
)
|
|
return
|
|
|
|
existing_error = (existing_running.error_message or "").strip()
|
|
stale_reason = (
|
|
f"Marked failed automatically after stale running timeout "
|
|
f"({RUNNING_TASK_GUARD_TIMEOUT_MINUTES}m) in scheduler guard"
|
|
)
|
|
existing_running.status = "failed"
|
|
existing_running.phase = "failed"
|
|
existing_running.completed_at = now
|
|
existing_running.error_message = (
|
|
f"{existing_error}\n{stale_reason}".strip()
|
|
if existing_error
|
|
else stale_reason
|
|
)
|
|
await db.commit()
|
|
logger.warning(
|
|
"Marked stale running task %s as failed before rerun of %s",
|
|
existing_running.id,
|
|
collector_name,
|
|
)
|
|
|
|
try:
|
|
collector._datasource_id = datasource.id
|
|
logger.info("Running collector: %s (datasource_id=%s)", collector_name, datasource.id)
|
|
task_result = await collector.run(db)
|
|
datasource.last_run_at = datetime.now(UTC)
|
|
datasource.last_status = task_result.get("status")
|
|
await _update_next_run_at(datasource, db)
|
|
logger.info("Collector %s completed: %s", collector_name, task_result)
|
|
except asyncio.CancelledError:
|
|
datasource.last_run_at = datetime.now(UTC)
|
|
datasource.last_status = "cancelled"
|
|
await db.commit()
|
|
logger.warning("Collector %s cancelled by operator", collector_name)
|
|
raise
|
|
except Exception as exc:
|
|
datasource.last_run_at = datetime.now(UTC)
|
|
datasource.last_status = "failed"
|
|
await db.commit()
|
|
logger.exception("Collector %s failed: %s", collector_name, exc)
|
|
|
|
|
|
async def cleanup_stale_running_tasks(max_age_hours: int = 2) -> int:
|
|
"""Mark stale running tasks as failed after restarts or collector hangs."""
|
|
cutoff = datetime.now(UTC) - timedelta(hours=max_age_hours)
|
|
|
|
async with async_session_factory() as db:
|
|
result = await db.execute(
|
|
select(CollectionTask).where(
|
|
CollectionTask.status == "running",
|
|
CollectionTask.started_at.is_not(None),
|
|
CollectionTask.started_at < cutoff,
|
|
)
|
|
)
|
|
stale_tasks = result.scalars().all()
|
|
|
|
for task in stale_tasks:
|
|
task.status = "failed"
|
|
task.phase = "failed"
|
|
task.completed_at = datetime.now(UTC)
|
|
existing_error = (task.error_message or "").strip()
|
|
cleanup_error = "Marked failed automatically after stale running task cleanup"
|
|
task.error_message = f"{existing_error}\n{cleanup_error}".strip() if existing_error else cleanup_error
|
|
|
|
if stale_tasks:
|
|
await db.commit()
|
|
logger.warning("Cleaned up %s stale running collection task(s)", len(stale_tasks))
|
|
|
|
return len(stale_tasks)
|
|
|
|
|
|
def start_scheduler() -> None:
|
|
"""Start the scheduler."""
|
|
if not scheduler.running:
|
|
scheduler.start()
|
|
logger.info("Scheduler started")
|
|
|
|
|
|
def stop_scheduler() -> None:
|
|
"""Stop the scheduler."""
|
|
if scheduler.running:
|
|
scheduler.shutdown(wait=False)
|
|
logger.info("Scheduler stopped")
|
|
|
|
|
|
async def sync_scheduler_with_datasources() -> None:
|
|
"""Synchronize scheduler jobs with datasource table."""
|
|
async with async_session_factory() as db:
|
|
result = await db.execute(select(DataSource).order_by(DataSource.id))
|
|
datasources = result.scalars().all()
|
|
|
|
configured_sources = {datasource.source for datasource in datasources}
|
|
for job in list(scheduler.get_jobs()):
|
|
if job.id not in configured_sources:
|
|
scheduler.remove_job(job.id)
|
|
|
|
for datasource in datasources:
|
|
await _apply_datasource_schedule(datasource, db)
|
|
|
|
|
|
async def sync_datasource_job(datasource_id: int) -> bool:
|
|
"""Synchronize a single datasource job after settings changes."""
|
|
async with async_session_factory() as db:
|
|
datasource = await db.get(DataSource, datasource_id)
|
|
if not datasource:
|
|
return False
|
|
|
|
await _apply_datasource_schedule(datasource, db)
|
|
return True
|
|
|
|
|
|
def get_scheduler_jobs() -> list[Dict[str, Any]]:
|
|
"""Get all scheduled jobs."""
|
|
jobs = []
|
|
for job in scheduler.get_jobs():
|
|
jobs.append(
|
|
{
|
|
"id": job.id,
|
|
"name": job.name,
|
|
"next_run_time": to_iso8601_utc(job.next_run_time),
|
|
"trigger": str(job.trigger),
|
|
}
|
|
)
|
|
return jobs
|
|
|
|
|
|
async def get_latest_task_id_for_datasource(datasource_id: int) -> Optional[int]:
|
|
from app.models.task import CollectionTask
|
|
|
|
async with async_session_factory() as db:
|
|
result = await db.execute(
|
|
select(CollectionTask.id)
|
|
.where(CollectionTask.datasource_id == datasource_id)
|
|
.order_by(CollectionTask.created_at.desc(), CollectionTask.id.desc())
|
|
.limit(1)
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
|
|
def run_collector_now(collector_name: str) -> bool:
|
|
"""Run a collector immediately (not scheduled)."""
|
|
collector = collector_registry.get(collector_name)
|
|
if not collector:
|
|
logger.error("Collector not found: %s", collector_name)
|
|
return False
|
|
|
|
try:
|
|
task = asyncio.create_task(run_collector_task(collector_name), name=_collector_task_name(collector_name))
|
|
RUNNING_COLLECTOR_TASKS[collector_name] = task
|
|
|
|
def _cleanup_task(done_task: asyncio.Task[Any]) -> None:
|
|
current = RUNNING_COLLECTOR_TASKS.get(collector_name)
|
|
if current is done_task:
|
|
RUNNING_COLLECTOR_TASKS.pop(collector_name, None)
|
|
|
|
task.add_done_callback(_cleanup_task)
|
|
logger.info("Triggered collector: %s", collector_name)
|
|
return True
|
|
except Exception as exc:
|
|
logger.error("Failed to trigger collector %s: %s", collector_name, exc)
|
|
return False
|
|
|
|
|
|
async def cancel_running_collector_now(collector_name: str) -> bool:
|
|
task = get_running_collector_task(collector_name)
|
|
if task is None or task.done():
|
|
RUNNING_COLLECTOR_TASKS.pop(collector_name, None)
|
|
return False
|
|
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
return True
|
|
return task.cancelled()
|