394 lines
14 KiB
Python
394 lines
14 KiB
Python
"""Task Scheduler for running collection jobs."""
|
|
|
|
import asyncio
|
|
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.core.logging import get_logger
|
|
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
|
|
from app.services.datasource_connectivity import (
|
|
build_builtin_connectivity_checksum,
|
|
get_builtin_effective_candidate,
|
|
save_connectivity_success,
|
|
)
|
|
|
|
logger = get_logger(__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_event(
|
|
"Collector not found for datasource",
|
|
event="collector.schedule.collector_missing",
|
|
context={"collector_name": 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_event(
|
|
"Scheduled collector",
|
|
event="collector.schedule.updated",
|
|
context={"collector_name": datasource.source, "frequency_minutes": datasource.frequency_minutes},
|
|
)
|
|
else:
|
|
logger.info_event(
|
|
"Collector disabled",
|
|
event="collector.schedule.disabled",
|
|
context={"collector_name": 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_event(
|
|
"Collector not found",
|
|
event="collector.run.collector_missing",
|
|
context={"collector_name": 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_event(
|
|
"Datasource not found for collector",
|
|
event="collector.run.datasource_missing",
|
|
context={"collector_name": collector_name},
|
|
)
|
|
return
|
|
|
|
if not datasource.is_active:
|
|
logger.info_event(
|
|
"Skipping disabled collector",
|
|
event="collector.run.skipped_disabled",
|
|
context={"collector_name": 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_event(
|
|
"Skipping collector trigger because task is already running",
|
|
event="collector.run.skipped_already_running",
|
|
context={"collector_name": collector_name, "task_id": 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_event(
|
|
"Marked stale running task as failed before rerun",
|
|
event="collector.run.stale_task_failed",
|
|
context={"collector_name": collector_name, "task_id": existing_running.id},
|
|
)
|
|
|
|
try:
|
|
collector._datasource_id = datasource.id
|
|
logger.info_event(
|
|
"Running collector",
|
|
event="collector.run.started",
|
|
context={"collector_name": collector_name, "datasource_id": datasource.id},
|
|
)
|
|
task_result = await collector.run(db)
|
|
datasource.last_run_at = datetime.now(UTC)
|
|
datasource.last_status = task_result.get("status")
|
|
if datasource.last_status == "success":
|
|
effective_candidate = await get_builtin_effective_candidate(db, datasource.source)
|
|
checksum, _credential_context = await build_builtin_connectivity_checksum(
|
|
datasource.source,
|
|
effective_candidate["endpoint"],
|
|
effective_candidate["auth_type"],
|
|
effective_candidate["headers"],
|
|
effective_candidate["config"],
|
|
db,
|
|
)
|
|
await save_connectivity_success(
|
|
db,
|
|
datasource.source,
|
|
checksum,
|
|
{"status_code": None},
|
|
connected_by="collection",
|
|
)
|
|
await _update_next_run_at(datasource, db)
|
|
logger.info_event(
|
|
"Collector completed",
|
|
event="collector.run.completed",
|
|
context={"collector_name": collector_name, "datasource_id": datasource.id, "result": task_result},
|
|
)
|
|
except asyncio.CancelledError:
|
|
datasource.last_run_at = datetime.now(UTC)
|
|
datasource.last_status = "cancelled"
|
|
await db.commit()
|
|
logger.warning_event(
|
|
"Collector cancelled by operator",
|
|
event="collector.run.cancelled",
|
|
context={"collector_name": collector_name, "datasource_id": datasource.id},
|
|
)
|
|
raise
|
|
except Exception as exc:
|
|
datasource.last_run_at = datetime.now(UTC)
|
|
datasource.last_status = "failed"
|
|
await db.commit()
|
|
logger.exception_event(
|
|
"Collector failed",
|
|
event="collector.run.failed",
|
|
context={"collector_name": collector_name, "datasource_id": datasource.id, "error": str(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_event(
|
|
"Cleaned up stale running collection tasks",
|
|
event="collector.cleanup.stale_tasks_cleaned",
|
|
context={"count": len(stale_tasks)},
|
|
)
|
|
|
|
return len(stale_tasks)
|
|
|
|
|
|
def start_scheduler() -> None:
|
|
"""Start the scheduler."""
|
|
if not scheduler.running:
|
|
scheduler.start()
|
|
logger.info_event("Scheduler started", event="scheduler.started")
|
|
|
|
|
|
def stop_scheduler() -> None:
|
|
"""Stop the scheduler."""
|
|
if scheduler.running:
|
|
scheduler.shutdown(wait=False)
|
|
logger.info_event("Scheduler stopped", event="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_event(
|
|
"Collector not found",
|
|
event="collector.trigger.collector_missing",
|
|
context={"collector_name": collector_name},
|
|
)
|
|
return False
|
|
|
|
existing_task = get_running_collector_task(collector_name)
|
|
if existing_task is not None and not existing_task.done():
|
|
logger.warning_event(
|
|
"Collector is already running in-memory; skipping duplicate trigger",
|
|
event="collector.trigger.skipped_already_running",
|
|
context={"collector_name": 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_event(
|
|
"Triggered collector",
|
|
event="collector.trigger.started",
|
|
context={"collector_name": collector_name},
|
|
)
|
|
return True
|
|
except Exception as exc:
|
|
logger.error_event(
|
|
"Failed to trigger collector",
|
|
event="collector.trigger.failed",
|
|
context={"collector_name": collector_name, "error": str(exc)},
|
|
)
|
|
return False
|
|
|
|
|
|
def is_collector_running(collector_name: str) -> bool:
|
|
task = get_running_collector_task(collector_name)
|
|
return bool(task is not None and not task.done())
|
|
|
|
|
|
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()
|