diff --git a/app/schemas/job.py b/app/schemas/job.py index 7ef5b8e..aee520e 100644 --- a/app/schemas/job.py +++ b/app/schemas/job.py @@ -152,7 +152,6 @@ class JobResponse(BaseModel): completed_at: datetime | None is_dlq: bool created_at: datetime - updated_at: datetime # Populated on detail endpoint only dependencies: list[UUID] | None = None diff --git a/app/services/dlq.py b/app/services/dlq.py index 8df85fc..66658bb 100644 --- a/app/services/dlq.py +++ b/app/services/dlq.py @@ -7,7 +7,7 @@ from app.core.logger import get_logger from app.models.job import Job, JobStatus from app.models.job_log import JobLog, LogEvent -from app.queues.base import BaseQueue +from app.queues.heapq import HeapQueue from app.services.settings import get_dlq_threshold logger = get_logger(__name__) @@ -16,12 +16,19 @@ async def send_to_dlq( job_id: uuid.UUID, error: str, - db: AsyncSession, + session: AsyncSession, ) -> None: """ Move a job to the dead letter queue after exhausting all retries. + + Steps: + 1. Set status='failed', is_dlq=True, last_error, worker_id=None + 2. Write job_failed log entry + 3. Cascade-cancel downstream DAG dependents + 4. Check DLQ count vs threshold — alert if breached + 5. Publish SSE event (handled by caller via redis publish) """ - await db.execute( + await session.execute( update(Job) .where(Job.id == job_id) .values( @@ -39,8 +46,8 @@ async def send_to_dlq( message=f"Job moved to DLQ after exhausting all retries. Error: {error}", metadata_={"error": error, "reason": "max_retries_exhausted"}, ) - db.add(log_entry) - await db.flush() + session.add(log_entry) + await session.flush() logger.error( "job_failed", @@ -49,16 +56,18 @@ async def send_to_dlq( reason="max_retries_exhausted", ) + # Cascade cancel downstream jobs in the DAG from app.services import dag - await dag.on_job_failed(job_id, db) + await dag.on_job_failed(job_id, session) - await _check_and_alert(db) + await _check_and_alert(session) async def _check_and_alert(session: AsyncSession) -> None: """ Count current DLQ jobs. If count >= threshold, fire alert email. + Imports alert_service lazily to avoid circular imports. """ count = await get_dlq_count(session) threshold = await get_dlq_threshold(session) @@ -71,29 +80,31 @@ async def _check_and_alert(session: AsyncSession) -> None: count=count, threshold=threshold, ) - from app.services import alert + from app.services import alert as alert_service - await alert.send_dlq_alert(count, session) + await alert_service.send_dlq_alert(count, session) async def get_dlq_jobs( page: int, limit: int, - db: AsyncSession, + session: AsyncSession, ) -> tuple[list[Job], int]: """ Return paginated DLQ jobs with total count. + Excludes soft-deleted jobs. + Ordered by most recently failed first. """ base_filter = [ Job.is_dlq.is_(True), Job.deleted_at.is_(None), ] - total_result = await db.execute(select(func.count(Job.id)).where(*base_filter)) + total_result = await session.execute(select(func.count(Job.id)).where(*base_filter)) total = total_result.scalar() or 0 offset = (page - 1) * limit - jobs_result = await db.execute( + jobs_result = await session.execute( select(Job) .where(*base_filter) .order_by(Job.updated_at.desc()) @@ -105,9 +116,9 @@ async def get_dlq_jobs( return jobs, total -async def get_dlq_count(db: AsyncSession) -> int: +async def get_dlq_count(session: AsyncSession) -> int: """Return the current number of jobs in the DLQ.""" - result = await db.execute( + result = await session.execute( select(func.count(Job.id)).where( Job.is_dlq.is_(True), Job.deleted_at.is_(None), @@ -117,13 +128,14 @@ async def get_dlq_count(db: AsyncSession) -> int: async def get_recent_dlq_jobs( - db: AsyncSession, + session: AsyncSession, limit: int = 10, ) -> list[dict]: """ Return the most recent DLQ jobs as plain dicts. + Used by alert_service to populate the email template. """ - result = await db.execute( + result = await session.execute( select(Job) .where( Job.is_dlq.is_(True), @@ -135,27 +147,35 @@ async def get_recent_dlq_jobs( jobs = result.scalars().all() return [ { - "id": str(job.id), - "short_id": str(job.id)[:8], - "type": job.type, - "last_error": job.last_error or "Unknown error", - "retry_count": job.retry_count, - "max_retries": job.max_retries, - "updated_at": job.updated_at.strftime("%Y-%m-%d %H:%M:%S UTC"), + "id": str(j.id), + "short_id": str(j.id)[:8], + "type": j.type, + "last_error": j.last_error or "Unknown error", + "retry_count": j.retry_count, + "max_retries": j.max_retries, + "updated_at": j.updated_at.strftime("%Y-%m-%d %H:%M:%S UTC"), } - for job in jobs + for j in jobs ] async def retry_dlq_job( job_id: uuid.UUID, - db: AsyncSession, - queue: BaseQueue, + session: AsyncSession, + queue: HeapQueue, ) -> Job: """ Manually retry a job from the DLQ. + + Steps: + 1. Verify job exists and is in DLQ + 2. Reset: status='pending', is_dlq=False, retry_count=0, + last_error=None, worker_id=None + 3. Call dag_service.on_dag_root_retried — reset downstream cancelled jobs + 4. Re-evaluate dependencies: if all met, push to queue + 5. Log and return updated job """ - result = await db.execute( + result = await session.execute( select(Job).where( Job.id == job_id, Job.deleted_at.is_(None), @@ -167,7 +187,7 @@ async def retry_dlq_job( raise JobNotInDLQException(str(job_id)) # Reset the job - await db.execute( + await session.execute( update(Job) .where(Job.id == job_id) .values( @@ -187,16 +207,16 @@ async def retry_dlq_job( message="Job manually retried from DLQ by engineer.", metadata_={"reason": "manual_dlq_retry"}, ) - db.add(log_entry) - await db.flush() + session.add(log_entry) + await session.flush() logger.info("dlq_job_retried", job_id=str(job_id)) - from app.services import dag + from app.services import dag as dag_service - await dag.on_dag_root_retried(job_id, db) + await dag_service.on_dag_root_retried(job_id, session) - has_unmet = await dag.has_unmet_dependencies(job_id, db) + has_unmet = await dag_service.has_unmet_dependencies(job_id, session) if not has_unmet: await queue.push( job_id=str(job_id), @@ -204,15 +224,11 @@ async def retry_dlq_job( scheduled_at=job.scheduled_at.timestamp(), created_at=job.created_at.timestamp(), ) - from app.services.job import _sync_job_to_redis - - await _sync_job_to_redis(str(job_id), float(job.priority)) - await db.commit() + await session.commit() - refreshed = await db.get(Job, job_id) - assert refreshed is not None - return refreshed + refreshed = await session.get(Job, job_id) + return refreshed # type: ignore async def remove_from_dlq( diff --git a/app/services/job.py b/app/services/job.py index 67bbf6e..b476fea 100644 --- a/app/services/job.py +++ b/app/services/job.py @@ -6,7 +6,7 @@ from sqlalchemy import func, or_, select, update from sqlalchemy.ext.asyncio import AsyncSession -from app.core.config import settings as app_settings +from app.core.config import settings from app.core.exceptions import ( JobNotCancellableException, JobNotDeletableException, @@ -17,77 +17,45 @@ from app.models.job import Job, JobStatus from app.models.job_log import JobLog, LogEvent from app.schemas.job import JobCreate, JobFilterParams, parse_interval -from app.services import dag +from app.services import dag as dag_service logger = get_logger(__name__) -QUEUE_KEY = "flint:queue" EVENTS_CHANNEL = "flint:events" -async def _get_redis() -> aioredis.Redis: - """Get a Redis client. Used internally by this service.""" - return aioredis.from_url( - app_settings.REDIS_URL, - decode_responses=True, - encoding="utf-8", - ) - - -async def _sync_job_to_redis(job_id: str, effective_priority: float) -> None: - """ - Add a job to the Redis sorted set (flint:queue). - Score = effective_priority. Lower score = higher urgency. - Workers read from this set to know which job to process next. - """ - redis = await _get_redis() - try: - await redis.zadd(QUEUE_KEY, {job_id: effective_priority}) - finally: - await redis.aclose() - - -async def _remove_job_from_redis(job_id: str) -> None: - """Remove a job from the Redis sorted set.""" - redis = await _get_redis() - try: - await redis.zrem(QUEUE_KEY, job_id) - finally: - await redis.aclose() - - async def _publish_sse_event(event: dict) -> None: - """ - Publish a job status change event to the Redis pub/sub channel. - The SSE endpoint subscribes to this channel and forwards events - to connected browser clients. - """ - redis = await _get_redis() + """Publish a job status change to the SSE Redis channel.""" try: + redis = aioredis.from_url( + settings.REDIS_URL, + decode_responses=True, + encoding="utf-8", + ) await redis.publish(EVENTS_CHANNEL, json.dumps(event)) - finally: await redis.aclose() + except Exception as exc: + logger.error("sse_publish_error", error=str(exc)) -async def create_job( - data: JobCreate, - db: AsyncSession, -) -> Job: +async def create_job(data: JobCreate, session: AsyncSession) -> Job: """ - Create a new job and optionally push it to the queue. + Create a new job and write it to PostgreSQL. + + The scheduler (worker/main.py) will pick it up on its next poll + and push it into the HeapQueue when scheduled_at becomes due + and all dependencies are met. """ - # Parse interval interval_seconds = None if data.interval: interval_seconds = parse_interval(data.interval) dependency_ids = data.dependency_ids or [] + + # Validate dependencies and check for cycles if dependency_ids: - await dag.check_cycle( - new_job_id=uuid.uuid4(), # placeholder — job not yet in DB - dependency_ids=dependency_ids, - db=db, - ) + temp_id = uuid.uuid4() + await dag_service.check_cycle(temp_id, dependency_ids, session) scheduled_at = data.scheduled_at or datetime.now(UTC) @@ -102,17 +70,13 @@ async def create_job( max_retries=data.max_retries, retry_count=0, ) - db.add(job) - await db.flush() + session.add(job) + await session.flush() - # Now run cycle check with the real job ID + # Re-run cycle check with the real job ID if dependency_ids: - await dag.check_cycle( - new_job_id=job.id, - dependency_ids=dependency_ids, - db=db, - ) - await dag.create_dependencies(job.id, dependency_ids, db) + await dag_service.check_cycle(job.id, dependency_ids, session) + await dag_service.create_dependencies(job.id, dependency_ids, session) log_entry = JobLog( job_id=job.id, @@ -126,8 +90,8 @@ async def create_job( "dependency_count": len(dependency_ids), }, ) - db.add(log_entry) - await db.commit() + session.add(log_entry) + await session.commit() logger.info( "job_created", @@ -138,12 +102,7 @@ async def create_job( has_dependencies=bool(dependency_ids), ) - has_unmet = bool(dependency_ids) - is_due = scheduled_at <= datetime.now(UTC) - - if not has_unmet and is_due: - await _sync_job_to_redis(str(job.id), job.effective_priority) - + # Notify SSE clients await _publish_sse_event( { "job_id": str(job.id), @@ -152,21 +111,25 @@ async def create_job( } ) + # Note: the scheduler will push this job into the HeapQueue + # on its next poll cycle when scheduled_at <= NOW and deps are met. + return job +# ------------------------------------------------------------------ # +# Read +# ------------------------------------------------------------------ # + + async def get_jobs( filters: JobFilterParams, - db: AsyncSession, + session: AsyncSession, ) -> tuple[list[Job], int]: - """ - Return paginated jobs with total count. - """ conditions: list = [ Job.deleted_at.is_(None), Job.is_dlq.is_(False), ] - if filters.status: conditions.append(Job.status == filters.status) if filters.type: @@ -174,51 +137,31 @@ async def get_jobs( if filters.priority: conditions.append(Job.priority == int(filters.priority)) if filters.search: - search_term = f"%{filters.search}%" + term = f"%{filters.search}%" conditions.append( or_( - Job.type.ilike(search_term), - Job.id.cast(db_string_type()).ilike(search_term), + Job.type.ilike(term), + Job.id.cast(type_=__import__("sqlalchemy").String).ilike(term), ) ) - total_result = await db.execute(select(func.count(Job.id)).where(*conditions)) + total_result = await session.execute(select(func.count(Job.id)).where(*conditions)) total = total_result.scalar() or 0 offset = (filters.page - 1) * filters.limit - jobs_result = await db.execute( + jobs_result = await session.execute( select(Job) .where(*conditions) .order_by(Job.created_at.desc()) .offset(offset) .limit(filters.limit) ) - jobs = list(jobs_result.scalars().all()) - - return jobs, total - + return list(jobs_result.scalars().all()), total -def db_string_type(): - """SQLAlchemy string type for UUID casting in search.""" - from sqlalchemy import String - return String - - -async def get_job_by_id( - job_id: uuid.UUID, - db: AsyncSession, - include_logs: bool = False, - include_dependencies: bool = False, -) -> Job: - """ - Fetch a single job by ID. - """ - result = await db.execute( - select(Job).where( - Job.id == job_id, - Job.deleted_at.is_(None), - ) +async def get_job_by_id(job_id: uuid.UUID, session: AsyncSession) -> Job: + result = await session.execute( + select(Job).where(Job.id == job_id, Job.deleted_at.is_(None)) ) job = result.scalar_one_or_none() if not job: @@ -228,42 +171,42 @@ async def get_job_by_id( async def get_job_with_details( job_id: uuid.UUID, - db: AsyncSession, + session: AsyncSession, ) -> tuple[Job, list[uuid.UUID], list]: - """ - Fetch a job with its dependency IDs and log entries. - """ - job = await get_job_by_id(job_id, db) - - dependency_ids = await dag.get_dependency_ids(job_id, db) + job = await get_job_by_id(job_id, session) + dependency_ids = await dag_service.get_dependency_ids(job_id, session) from sqlalchemy import asc - from app.models.job_log import JobLog + from app.models.job_log import JobLog as JobLogModel - logs_result = await db.execute( - select(JobLog).where(JobLog.job_id == job_id).order_by(asc(JobLog.created_at)) + logs_result = await session.execute( + select(JobLogModel) + .where(JobLogModel.job_id == job_id) + .order_by(asc(JobLogModel.created_at)) ) logs = list(logs_result.scalars().all()) - return job, dependency_ids, logs -async def cancel_job( - job_id: uuid.UUID, - db: AsyncSession, -) -> Job: +# ------------------------------------------------------------------ # +# Cancel +# ------------------------------------------------------------------ # + + +async def cancel_job(job_id: uuid.UUID, session: AsyncSession) -> Job: """ - Request cancellation of a job. + Pending jobs: mark cancelled immediately. + Processing jobs: set cancellation_requested=True (cooperative). + Terminal jobs: raise JobNotCancellableException. """ - job = await get_job_by_id(job_id, db) + job = await get_job_by_id(job_id, session) if not job.is_cancellable: raise JobNotCancellableException(str(job_id), job.status) if job.status == JobStatus.PENDING: - # Immediate cancellation - await db.execute( + await session.execute( update(Job) .where(Job.id == job_id) .values( @@ -272,180 +215,127 @@ async def cancel_job( updated_at=func.now(), ) ) - - await _remove_job_from_redis(str(job_id)) - - # Cascade cancel downstream dependents - await dag.on_job_failed(job_id, db) - - log_entry = JobLog( - job_id=job_id, - event=LogEvent.JOB_CANCELLED, - message="Job cancelled by user request while pending.", - metadata_={"reason": "user_requested", "was_status": "pending"}, + await dag_service.on_job_failed(job_id, session) + + session.add( + JobLog( + job_id=job_id, + event=LogEvent.JOB_CANCELLED, + message="Job cancelled while pending.", + metadata_={"reason": "user_requested", "was_status": "pending"}, + ) ) - db.add(log_entry) - logger.info("job_cancelled", job_id=str(job_id), was_status="pending") elif job.status == JobStatus.PROCESSING: - await db.execute( + await session.execute( update(Job) .where(Job.id == job_id) - .values( - cancellation_requested=True, - updated_at=func.now(), - ) + .values(cancellation_requested=True, updated_at=func.now()) ) - - log_entry = JobLog( - job_id=job_id, - event=LogEvent.JOB_CANCELLED, - message=( - "Cancellation requested while job is processing. " - "Worker will honour at next checkpoint." - ), - metadata_={"reason": "user_requested", "was_status": "processing"}, + session.add( + JobLog( + job_id=job_id, + event=LogEvent.JOB_CANCELLED, + message=( + "Cancellation requested while processing. " + "Worker will honour at next checkpoint." + ), + metadata_={"reason": "user_requested", "was_status": "processing"}, + ) ) - db.add(log_entry) - logger.info( "job_cancellation_requested", job_id=str(job_id), was_status="processing", ) - await db.commit() - - await _publish_sse_event( - { - "job_id": str(job_id), - "status": JobStatus.CANCELLED, - } - ) - - job = await db.get(Job, job_id) + await session.commit() + await _publish_sse_event({"job_id": str(job_id), "status": JobStatus.CANCELLED}) + job = await session.get(Job, job_id) assert job is not None, f"Job {job_id} vanished after cancellation" return job -async def soft_delete_job( - job_id: uuid.UUID, - db: AsyncSession, -) -> None: - """ - Move a job to the bin (sets deleted_at). - """ - job = await get_job_by_id(job_id, db) +# ------------------------------------------------------------------ # +# Soft delete / Bin / Restore / Hard delete +# ------------------------------------------------------------------ # + +async def soft_delete_job(job_id: uuid.UUID, session: AsyncSession) -> None: + job = await get_job_by_id(job_id, session) if not job.is_terminal: raise JobNotDeletableException(str(job_id), job.status) - - await db.execute( + await session.execute( update(Job) .where(Job.id == job_id) .values(deleted_at=func.now(), updated_at=func.now()) ) - await db.commit() - + await session.commit() logger.info("job_soft_deleted", job_id=str(job_id)) async def get_bin_jobs( - page: int, - limit: int, - db: AsyncSession, + page: int, limit: int, session: AsyncSession ) -> tuple[list[Job], int]: - """ - Return soft-deleted jobs (the bin). Paginated, most recently deleted first. - """ conditions = [Job.deleted_at.isnot(None)] - - total_result = await db.execute(select(func.count(Job.id)).where(*conditions)) - total = total_result.scalar() or 0 - + total = ( + await session.execute(select(func.count(Job.id)).where(*conditions)) + ).scalar() or 0 offset = (page - 1) * limit - jobs_result = await db.execute( - select(Job) - .where(*conditions) - .order_by(Job.deleted_at.desc()) - .offset(offset) - .limit(limit) + jobs = list( + ( + await session.execute( + select(Job) + .where(*conditions) + .order_by(Job.deleted_at.desc()) + .offset(offset) + .limit(limit) + ) + ) + .scalars() + .all() ) - jobs = list(jobs_result.scalars().all()) - return jobs, total -async def restore_job( - job_id: uuid.UUID, - db: AsyncSession, -) -> Job: - """ - Restore a soft-deleted job from the bin. - Clears deleted_at. Does not re-queue the job. - Raises JobNotInBinException if the job is not soft-deleted. - """ - result = await db.execute(select(Job).where(Job.id == job_id)) +async def restore_job(job_id: uuid.UUID, session: AsyncSession) -> Job: + result = await session.execute(select(Job).where(Job.id == job_id)) job = result.scalar_one_or_none() - if not job or job.deleted_at is None: raise JobNotInBinException(str(job_id)) - - await db.execute( + await session.execute( update(Job) .where(Job.id == job_id) .values(deleted_at=None, updated_at=func.now()) ) - await db.commit() - + await session.commit() logger.info("job_restored", job_id=str(job_id)) - job = await db.get(Job, job_id) + job = await session.get(Job, job_id) assert job is not None, f"Job {job_id} vanished after cancellation" return job -async def hard_delete_job( - job_id: uuid.UUID, - db: AsyncSession, -) -> None: - """ - Permanently delete a job from the database. - """ - result = await db.execute(select(Job).where(Job.id == job_id)) +async def hard_delete_job(job_id: uuid.UUID, session: AsyncSession) -> None: + result = await session.execute(select(Job).where(Job.id == job_id)) job = result.scalar_one_or_none() - if not job or job.deleted_at is None: raise JobNotInBinException(str(job_id)) - - await db.delete(job) - await db.commit() - + await session.delete(job) + await session.commit() logger.info("job_hard_deleted", job_id=str(job_id)) -async def get_job_counts_by_status( - db: AsyncSession, -) -> dict[str, int]: - """ - Return a dict of job counts per status for the dashboard. - Excludes soft-deleted jobs. Includes DLQ count separately. - """ - +async def get_job_counts_by_status(session: AsyncSession) -> dict[str, int]: from app.services.dlq import get_dlq_count - results = await db.execute( + results = await session.execute( select(Job.status, func.count(Job.id)) - .where( - Job.deleted_at.is_(None), - Job.is_dlq.is_(False), - ) + .where(Job.deleted_at.is_(None), Job.is_dlq.is_(False)) .group_by(Job.status) ) counts = {status: count for status, count in results.fetchall()} - - dlq_count = await get_dlq_count(db) - + dlq_count = await get_dlq_count(session) return { "pending": counts.get(JobStatus.PENDING, 0), "processing": counts.get(JobStatus.PROCESSING, 0), diff --git a/docker-compose.yml b/docker-compose.yml index f597285..95e1216 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -43,12 +43,12 @@ services: condition: service_healthy restart: unless-stopped - worker-1: + worker: build: . - command: uv run python -m worker.worker + command: uv run python -m worker.main env_file: .env environment: - WORKER_ID: worker-1 + WORKER_COUNT: 2 DATABASE_URL: postgresql+asyncpg://${POSTGRES_USER:-flint}:${POSTGRES_PASSWORD}@postgres:5432/${POSTGRES_DB:-flint} REDIS_URL: redis://redis:6379/0 depends_on: @@ -59,39 +59,6 @@ services: volumes: - ./logs:/app/logs - worker-2: - build: . - command: uv run python -m worker.worker - env_file: .env - environment: - WORKER_ID: worker-2 - DATABASE_URL: postgresql+asyncpg://${POSTGRES_USER:-flint}:${POSTGRES_PASSWORD}@postgres:5432/${POSTGRES_DB:-flint} - REDIS_URL: redis://redis:6379/0 - depends_on: - postgres: - condition: service_healthy - redis: - condition: service_healthy - volumes: - - ./logs:/app/logs - - scheduler: - build: . - command: uv run python -m scheduler.scheduler - env_file: .env - environment: - DATABASE_URL: postgresql+asyncpg://${POSTGRES_USER:-flint}:${POSTGRES_PASSWORD}@postgres:5432/${POSTGRES_DB:-flint} - REDIS_URL: redis://redis:6379/0 - depends_on: - postgres: - condition: service_healthy - redis: - condition: service_healthy - volumes: - - ./logs:/app/logs - - - volumes: postgres_data: redis_data: diff --git a/scheduler/scheduler.py b/scheduler/scheduler.py index dd7e543..497b4a4 100644 --- a/scheduler/scheduler.py +++ b/scheduler/scheduler.py @@ -1,72 +1,50 @@ import asyncio -import signal from datetime import UTC, datetime -import redis.asyncio as aioredis from sqlalchemy import func, select from app.core.config import settings -from app.core.logger import get_logger, setup_logging +from app.core.logger import get_logger from app.db.session import AsyncSessionLocal from app.models.job import Job, JobStatus from app.models.job_depedencies import JobDependency from app.queues.heapq import HeapQueue -from worker.aging import AgingProcess -setup_logging() logger = get_logger(__name__) -QUEUE_KEY = "flint:queue" -WORKER_KEY_PATTERN = "flint:worker:*" -DEAD_WORKER_CLEANUP_INTERVAL = 60 # seconds - class FlintScheduler: - def __init__(self) -> None: - self.running = False - self._redis: aioredis.Redis - self._queue = HeapQueue() - - async def start(self) -> None: - """Start all scheduler loops concurrently.""" - self.running = True - - self._redis = aioredis.from_url( - settings.REDIS_URL, - decode_responses=True, - encoding="utf-8", - ) - - loop = asyncio.get_event_loop() - for sig in (signal.SIGTERM, signal.SIGINT): - loop.add_signal_handler(sig, self._handle_shutdown) + def __init__(self, queue: HeapQueue) -> None: + self.queue = queue + async def run(self, shutdown_event: asyncio.Event) -> None: + """ + Main scheduler coroutine. Runs until shutdown_event is set. + On startup, loads all currently pending due jobs into the queue. + """ logger.info("scheduler_started") - await asyncio.gather( - self._due_job_loop(), - self._aging_loop(), - self._dead_worker_cleanup_loop(), - ) + await self._push_due_jobs() - async def _due_job_loop(self) -> None: - """ - Every SCHEDULER_POLL_INTERVAL seconds: - Find all pending jobs that are due and eligible, push to Redis queue. - """ - while self.running: + while not shutdown_event.is_set(): try: await self._push_due_jobs() except asyncio.CancelledError: break except Exception as exc: - logger.error("scheduler_due_job_error", error=str(exc)) - await asyncio.sleep(settings.SCHEDULER_POLL_INTERVAL) + logger.error("scheduler_error", error=str(exc)) + + try: + await asyncio.wait_for( + shutdown_event.wait(), + timeout=settings.SCHEDULER_POLL_INTERVAL, + ) + except TimeoutError: + pass + + logger.info("scheduler_stopped") async def _push_due_jobs(self) -> None: - """ - Query for eligible jobs and push them to the Redis sorted set. - """ async with AsyncSessionLocal() as session: now = datetime.now(UTC) @@ -84,26 +62,26 @@ async def _push_due_jobs(self) -> None: Job.scheduled_at.asc(), Job.created_at.asc(), ) - .limit(100) + .limit(200) ) - candidate_jobs = result.scalars().all() + candidates = result.scalars().all() - if not candidate_jobs: + if not candidates: return pushed = 0 - for job in candidate_jobs: - score = await self._redis.zscore(QUEUE_KEY, str(job.id)) - if score is not None: + for job in candidates: + if await self.queue.contains(str(job.id)): continue - unmet = await self._has_unmet_dependencies(job.id, session) - if unmet: + if await self._has_unmet_dependencies(job.id, session): continue - await self._redis.zadd( - QUEUE_KEY, - {str(job.id): job.effective_priority}, + await self.queue.push( + job_id=str(job.id), + effective_priority=job.effective_priority, + scheduled_at=job.scheduled_at.timestamp(), + created_at=job.created_at.timestamp(), ) pushed += 1 @@ -111,15 +89,11 @@ async def _push_due_jobs(self) -> None: logger.info( "jobs_pushed_to_queue", count=pushed, - timestamp=now.isoformat(), + queue_size=await self.queue.size(), ) - async def _has_unmet_dependencies( - self, - job_id, - session, - ) -> bool: - """Return True if the job has any dependency that is not completed.""" + async def _has_unmet_dependencies(self, job_id, session) -> bool: + """Return True if the job has any dependency not yet completed.""" result = await session.execute( select(func.count(JobDependency.id)) .join(Job, Job.id == JobDependency.depends_on_id) @@ -129,107 +103,3 @@ async def _has_unmet_dependencies( ) ) return (result.scalar() or 0) > 0 - - async def _aging_loop(self) -> None: - """ - Every AGING_INTERVAL seconds: run the starvation prevention - aging process to decrement effective_priority on waiting jobs. - """ - while self.running: - try: - async with AsyncSessionLocal() as session: - aging = AgingProcess() - await aging.run(session, self._queue) - except asyncio.CancelledError: - break - except Exception as exc: - logger.error("scheduler_aging_error", error=str(exc)) - await asyncio.sleep(settings.AGING_INTERVAL) - - async def _dead_worker_cleanup_loop(self) -> None: - """ - Every 60 seconds: find jobs stuck in 'processing' state whose - worker heartbeat has expired in Redis. Reset them to 'pending' - so they get picked up again. - """ - while self.running: - try: - await self._cleanup_dead_worker_jobs() - except asyncio.CancelledError: - break - except Exception as exc: - logger.error("scheduler_cleanup_error", error=str(exc)) - await asyncio.sleep(DEAD_WORKER_CLEANUP_INTERVAL) - - async def _cleanup_dead_worker_jobs(self) -> None: - """ - Find active worker IDs from Redis. Any job with a worker_id - not in that set has been abandoned by a dead worker and should - be reset to pending. - """ - worker_keys = await self._redis.keys(WORKER_KEY_PATTERN) - active_worker_ids = {key.replace("flint:worker:", "") for key in worker_keys} # type: ignore - - async with AsyncSessionLocal() as session: - result = await session.execute( - select(Job).where( - Job.status == JobStatus.PROCESSING, - Job.deleted_at.is_(None), - Job.worker_id.isnot(None), - ) - ) - processing_jobs = result.scalars().all() - - reset_count = 0 - for job in processing_jobs: - if job.worker_id not in active_worker_ids: - from sqlalchemy import update - - await session.execute( - update(Job) - .where(Job.id == job.id) - .values( - status=JobStatus.PENDING, - worker_id=None, - updated_at=func.now(), - ) - ) - - from app.models.job_log import JobLog, LogEvent - - log = JobLog( - job_id=job.id, - event=LogEvent.JOB_CREATED, - message=( - f"Job reset to pending: worker '{job.worker_id}' " - f"is no longer active (heartbeat expired)." - ), - metadata_={ - "dead_worker_id": job.worker_id, - "reason": "dead_worker_recovery", - }, - ) - session.add(log) - reset_count += 1 - - logger.warning( - "dead_worker_job_reset", - job_id=str(job.id), - dead_worker_id=job.worker_id, - ) - - if reset_count > 0: - await session.commit() - logger.info( - "dead_worker_cleanup_complete", - reset_count=reset_count, - ) - - def _handle_shutdown(self) -> None: - logger.info("scheduler_shutdown_signal") - self.running = False - - -if __name__ == "__main__": - scheduler = FlintScheduler() - asyncio.run(scheduler.start()) diff --git a/tests/conftest.py b/tests/conftest.py deleted file mode 100644 index af7e479..0000000 --- a/tests/conftest.py +++ /dev/null @@ -1,6 +0,0 @@ -import pytest - - -@pytest.fixture -def anyio_backend(): - return "asyncio" diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py deleted file mode 100644 index c210fac..0000000 --- a/tests/integration/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Integration tests.""" diff --git a/worker/aging.py b/worker/aging.py index e31c953..504534f 100644 --- a/worker/aging.py +++ b/worker/aging.py @@ -1,16 +1,12 @@ from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING -from sqlalchemy import func, select, update +from sqlalchemy import Float, cast, func, select, update from sqlalchemy.ext.asyncio import AsyncSession from app.core.config import settings from app.core.logger import get_logger from app.models.job import Job, JobPriority, JobStatus - -if TYPE_CHECKING: - from app.queues.base import BaseQueue - +from app.queues.heapq import HeapQueue logger = get_logger(__name__) @@ -18,39 +14,26 @@ class AgingProcess: - """ - Decrements effective_priority on long-waiting pending jobs. - """ - - async def run( - self, - session: AsyncSession, - queue: "BaseQueue", - ) -> dict: + async def run(self, session: AsyncSession, queue: HeapQueue) -> dict: """ - Execute one aging cycle. + One aging cycle: + 1. Decrement effective_priority for qualifying medium + low jobs in DB + 2. For each updated job, call queue.update_priority() on the heap + 3. Log and return summary """ now = datetime.now(UTC) medium_cutoff = now - timedelta(seconds=settings.MEDIUM_PRIORITY_AGE_THRESHOLD) low_cutoff = now - timedelta(seconds=settings.LOW_PRIORITY_AGE_THRESHOLD) - medium_count = await self._age_priority( - priority=JobPriority.MEDIUM, - cutoff=medium_cutoff, - db=session, - ) - - low_count = await self._age_priority( - priority=JobPriority.LOW, - cutoff=low_cutoff, - db=session, - ) + medium_count = await self._age(JobPriority.MEDIUM, medium_cutoff, session) + low_count = await self._age(JobPriority.LOW, low_cutoff, session) await session.commit() + # Sync updated priorities into the heap if medium_count > 0 or low_count > 0: - await self._sync_queue(session, queue) + await self._sync_heap(session, queue) summary = { "medium_jobs_aged": medium_count, @@ -58,28 +41,16 @@ async def run( "total_aged": medium_count + low_count, "timestamp": now.isoformat(), } - - logger.info( - "aging_complete", - **summary, - ) - + logger.info("aging_complete", **summary) return summary - async def _age_priority( - self, - priority: int, - cutoff: datetime, - db: AsyncSession, - ) -> int: + async def _age(self, priority: int, cutoff: datetime, session: AsyncSession) -> int: """ - Decrement effective_priority for all pending jobs of a given - priority level that have been waiting since before the cutoff. - - Uses GREATEST() to floor at _PRIORITY_FLOOR (1.0). - Returns the number of rows updated. + Decrement effective_priority for pending jobs of a given priority + that have been waiting since before the cutoff. + Returns number of rows updated. """ - result = await db.execute( + result = await session.execute( update(Job) .where( Job.status == JobStatus.PENDING, @@ -98,56 +69,36 @@ async def _age_priority( ) .returning(Job.id) ) - updated_ids = result.fetchall() - count = len(updated_ids) - + updated = result.fetchall() + count = len(updated) if count > 0: - logger.info( - "jobs_aged", - priority=priority, - count=count, - decrement=settings.AGING_DECREMENT, - ) - + logger.info("jobs_aged", priority=priority, count=count) return count - async def _sync_queue( - self, - db: AsyncSession, - queue: "BaseQueue", - ) -> None: + async def _sync_heap(self, session: AsyncSession, queue: HeapQueue) -> None: """ - After aging, fetch all updated pending jobs and re-push them - to the queue with their new effective_priority scores. + After aging, reload all pending jobs whose effective_priority + has dropped below their raw priority and update them in the heap. """ - result = await db.execute( + result = await session.execute( select(Job).where( Job.status == JobStatus.PENDING, Job.deleted_at.is_(None), Job.is_dlq.is_(False), - # Only re-sync jobs that have been aged (below their raw priority) - Job.effective_priority - < Job.priority.cast(type_=type(Job.effective_priority.type)), + Job.effective_priority < cast(Job.priority, Float), ) ) aged_jobs = result.scalars().all() for job in aged_jobs: try: - await queue.update_priority( - job_id=str(job.id), - new_priority=job.effective_priority, - scheduled_at=job.scheduled_at.timestamp(), - created_at=job.created_at.timestamp(), - ) - - # Also update Redis sorted set score - from app.services.job import _sync_job_to_redis - - await _sync_job_to_redis( - str(job.id), - job.effective_priority, - ) + if await queue.contains(str(job.id)): + await queue.update_priority( + job_id=str(job.id), + new_priority=job.effective_priority, + scheduled_at=job.scheduled_at.timestamp(), + created_at=job.created_at.timestamp(), + ) except Exception as exc: logger.error( "aging_sync_error", diff --git a/worker/main.py b/worker/main.py new file mode 100644 index 0000000..5b09f91 --- /dev/null +++ b/worker/main.py @@ -0,0 +1,98 @@ +import asyncio +import os +import signal + +from app.core.logger import get_logger, setup_logging +from app.queues.heapq import HeapQueue +from scheduler.scheduler import FlintScheduler +from worker.aging import AgingProcess +from worker.worker import run_worker + +setup_logging() +logger = get_logger(__name__) + + +async def main() -> None: + worker_count = int(os.getenv("WORKER_COUNT", "2")) + queue = HeapQueue() + + logger.info( + "flint_starting", + worker_count=worker_count, + queue_type="HeapQueue", + ) + + scheduler = FlintScheduler(queue=queue) + aging = AgingProcess() + + shutdown_event = asyncio.Event() + + def _handle_signal(): + logger.info("shutdown_signal_received") + shutdown_event.set() + + loop = asyncio.get_event_loop() + for sig in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(sig, _handle_signal) + + tasks = [ + asyncio.create_task(scheduler.run(shutdown_event), name="scheduler"), + asyncio.create_task( + _aging_loop(aging, queue, shutdown_event), + name="aging", + ), + ] + + for i in range(worker_count): + tasks.append( + asyncio.create_task( + run_worker( + worker_id=f"worker-{i + 1}", + queue=queue, + shutdown_event=shutdown_event, + ), + name=f"worker-{i + 1}", + ) + ) + + logger.info("all_tasks_started", task_count=len(tasks)) + + await shutdown_event.wait() + + logger.info("shutting_down") + + for task in tasks: + task.cancel() + + await asyncio.gather(*tasks, return_exceptions=True) + logger.info("flint_stopped") + + +async def _aging_loop( + aging: AgingProcess, + queue: HeapQueue, + shutdown_event: asyncio.Event, +) -> None: + from app.core.config import settings + from app.db.session import AsyncSessionLocal + + while not shutdown_event.is_set(): + try: + async with AsyncSessionLocal() as session: + await aging.run(session, queue) + except asyncio.CancelledError: + break + except Exception as exc: + logger.error("aging_loop_error", error=str(exc)) + + try: + await asyncio.wait_for( + shutdown_event.wait(), + timeout=settings.AGING_INTERVAL, + ) + except TimeoutError: + pass + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/worker/processor.py b/worker/processor.py index 16d1c3f..ae9f4a9 100644 --- a/worker/processor.py +++ b/worker/processor.py @@ -3,7 +3,6 @@ import random import uuid from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING from sqlalchemy import func, select, update from sqlalchemy.ext.asyncio import AsyncSession @@ -13,49 +12,57 @@ from app.handlers import get_handler from app.models.job import Job, JobStatus from app.models.job_log import JobLog, LogEvent -from app.services import dag as dag_service +from app.queues.heapq import HeapQueue +from app.services import dag from app.services.dlq import send_to_dlq -from app.services.job import ( - _publish_sse_event, - _sync_job_to_redis, -) - -if TYPE_CHECKING: - from app.queues.base import BaseQueue logger = get_logger(__name__) +_EVENTS_CHANNEL = "flint:events" + def calculate_next_retry_delay(attempt: int) -> float: """ - Exponential backoff with jitter. + Exponential backoff with full jitter. """ base = math.pow(5, attempt - 1) - jitter = base * random.uniform(0.5, 1.5) - return round(jitter, 2) + return round(base * random.uniform(0.5, 1.5), 2) + + +async def _publish_sse(event: dict) -> None: + """Publish a job status event to Redis pub/sub for SSE streaming.""" + import json + + import redis.asyncio as aioredis + + from app.core.config import settings + + try: + redis = aioredis.from_url(settings.REDIS_URL, decode_responses=True) + await redis.publish(_EVENTS_CHANNEL, json.dumps(event)) + await redis.aclose() + except Exception as exc: + logger.error("sse_publish_error", error=str(exc)) class JobProcessor: - def __init__(self, worker_id: str, queue: "BaseQueue") -> None: + def __init__(self, worker_id: str, queue: HeapQueue) -> None: self.worker_id = worker_id self.queue = queue async def process(self, job_id: str) -> None: - """ - Full job processing flow for a single job. - """ - async with AsyncSessionLocal() as db: - job = await self._load_job(job_id, db) + """Full processing lifecycle for a single job.""" + async with AsyncSessionLocal() as session: + job = await self._load_job(job_id, session) if not job: return - claimed = await self._claim_job(job.id, db) + claimed = await self._claim_job(job.id, session) if not claimed: logger.info( "job_claim_failed", job_id=job_id, worker_id=self.worker_id, - reason="already_claimed", ) return @@ -64,9 +71,9 @@ async def process(self, job_id: str) -> None: event=LogEvent.JOB_STARTED, message=f"Job started by worker {self.worker_id}.", metadata={"worker_id": self.worker_id}, - db=db, + session=session, ) - await _publish_sse_event( + await _publish_sse( { "job_id": job_id, "status": JobStatus.PROCESSING, @@ -82,11 +89,9 @@ async def process(self, job_id: str) -> None: retry_count=job.retry_count, ) - if await self._is_cancellation_requested(job.id, db): + if await self._is_cancelled(job.id, session): await self._mark_cancelled( - job.id, - db, - reason="cancellation_requested_before_execute", + job.id, session, "cancellation_requested_before_execute" ) return @@ -95,29 +100,26 @@ async def process(self, job_id: str) -> None: handler = get_handler(job.type) result = await handler.execute(job.payload) except Exception as exc: - elapsed_ms = self._elapsed_ms(start_time) + elapsed = self._elapsed_ms(start_time) logger.warning( "job_execution_failed", job_id=job_id, - type=job.type, error=str(exc), - elapsed_ms=elapsed_ms, + elapsed_ms=elapsed, retry_count=job.retry_count, ) - await self._handle_failure(job, exc, db) + await self._handle_failure(job, exc, session) return - elapsed_ms = self._elapsed_ms(start_time) + elapsed = self._elapsed_ms(start_time) - if await self._is_cancellation_requested(job.id, db): + if await self._is_cancelled(job.id, session): await self._mark_cancelled( - job.id, - db, - reason="cancellation_requested_after_execute", + job.id, session, "cancellation_requested_after_execute" ) return - await db.execute( + await session.execute( update(Job) .where(Job.id == job.id) .values( @@ -131,56 +133,45 @@ async def process(self, job_id: str) -> None: await self._log_event( job_id=job.id, event=LogEvent.JOB_COMPLETED, - message=( - f"Job completed successfully by worker {self.worker_id} " - f"in {elapsed_ms}ms." - ), + message=f"Job completed by worker {self.worker_id} in {elapsed}ms.", metadata={ "worker_id": self.worker_id, - "duration_ms": elapsed_ms, + "duration_ms": elapsed, "result": result, }, - db=db, + session=session, ) - await db.commit() + await session.commit() logger.info( "job_completed", job_id=job_id, type=job.type, worker_id=self.worker_id, - duration_ms=elapsed_ms, + duration_ms=elapsed, ) - await _publish_sse_event( + await _publish_sse( { "job_id": job_id, "status": JobStatus.COMPLETED, "worker_id": self.worker_id, - "duration_ms": elapsed_ms, + "duration_ms": elapsed, } ) - await dag_service.on_job_completed(job.id, db, self.queue) + await dag.on_job_completed(job.id, session, self.queue) if job.interval_seconds: - await self._handle_recurrence(job, db) + await self._handle_recurrence(job, session) - async def _claim_job( - self, - job_id: uuid.UUID, - db: AsyncSession, - ) -> bool: + async def _claim_job(self, job_id: uuid.UUID, session: AsyncSession) -> bool: """ - Atomic claim via single UPDATE with WHERE guard. - - Only succeeds if status='pending' AND worker_id IS NULL. - If two workers race, exactly one gets the row back. - PostgreSQL row-level locking guarantees atomicity. - - Returns True if this worker successfully claimed the job. + Atomically claim a job. + Returns True only if this worker successfully set worker_id. + PostgreSQL row-level atomicity ensures only one winner. """ - result = await db.execute( + result = await session.execute( update(Job) .where( Job.id == job_id, @@ -196,18 +187,15 @@ async def _claim_job( ) .returning(Job.id) ) - await db.commit() + await session.commit() return result.scalar_one_or_none() is not None async def _handle_failure( self, job: Job, error: Exception, - db: AsyncSession, + session: AsyncSession, ) -> None: - """ - Handle a job execution failure. - """ new_retry_count = job.retry_count + 1 error_str = str(error) @@ -215,7 +203,7 @@ async def _handle_failure( delay = calculate_next_retry_delay(new_retry_count) next_retry_at = datetime.now(UTC) + timedelta(seconds=delay) - await db.execute( + await session.execute( update(Job) .where(Job.id == job.id) .values( @@ -232,19 +220,18 @@ async def _handle_failure( job_id=job.id, event=LogEvent.JOB_RETRY_ATTEMPTED, message=( - f"Job failed on attempt {new_retry_count}/{job.max_retries}. " + f"Attempt {new_retry_count}/{job.max_retries} failed. " f"Retrying in {delay:.1f}s. Error: {error_str[:200]}" ), metadata={ "attempt": new_retry_count, - "max_retries": job.max_retries, "delay_seconds": delay, "error": error_str[:500], "next_retry_at": next_retry_at.isoformat(), }, - db=db, + session=session, ) - await db.commit() + await session.commit() logger.warning( "job_retry_attempted", @@ -252,10 +239,9 @@ async def _handle_failure( attempt=new_retry_count, max_retries=job.max_retries, delay_seconds=delay, - error=error_str[:200], ) - await _publish_sse_event( + await _publish_sse( { "job_id": str(job.id), "status": JobStatus.PENDING, @@ -263,30 +249,20 @@ async def _handle_failure( } ) - asyncio.create_task( - self._retry_after_delay( - job_id=str(job.id), - effective_priority=job.effective_priority, - scheduled_at=next_retry_at.timestamp(), - created_at=job.created_at.timestamp(), - delay=delay, - ) - ) + asyncio.create_task(self._requeue_after_delay(job, next_retry_at, delay)) else: - await db.execute( + await session.execute( update(Job) .where(Job.id == job.id) .values( - retry_count=new_retry_count, - worker_id=None, - updated_at=func.now(), + retry_count=new_retry_count, worker_id=None, updated_at=func.now() ) ) - await db.flush() - await send_to_dlq(job.id, error_str[:1000], db) + await session.flush() + await send_to_dlq(job.id, error_str[:1000], session) - await _publish_sse_event( + await _publish_sse( { "job_id": str(job.id), "status": JobStatus.FAILED, @@ -294,45 +270,31 @@ async def _handle_failure( } ) - async def _retry_after_delay( + async def _requeue_after_delay( self, - job_id: str, - effective_priority: float, - scheduled_at: float, - created_at: float, + job: Job, + next_retry_at: datetime, delay: float, ) -> None: """ - Background task: wait for the backoff delay then push the - job back onto the queue so it gets picked up again. + Wait for the backoff delay then push the job back into the heap. + The scheduler would also pick it up on its next poll — whichever + happens first, the atomic claim prevents double-processing. """ await asyncio.sleep(delay) try: await self.queue.push( - job_id=job_id, - effective_priority=effective_priority, - scheduled_at=scheduled_at, - created_at=created_at, + job_id=str(job.id), + effective_priority=job.effective_priority, + scheduled_at=next_retry_at.timestamp(), + created_at=job.created_at.timestamp(), ) - await _sync_job_to_redis(job_id, effective_priority) - logger.info("job_retry_queued", job_id=job_id, delay=delay) + logger.info("job_retry_requeued", job_id=str(job.id), delay=delay) except Exception as exc: - logger.error( - "job_retry_queue_error", - job_id=job_id, - error=str(exc), - ) + logger.error("job_retry_requeue_error", job_id=str(job.id), error=str(exc)) - async def _is_cancellation_requested( - self, - job_id: uuid.UUID, - db: AsyncSession, - ) -> bool: - """ - Check the cancellation_requested flag from the DB. - Called at checkpoints during processing. - """ - result = await db.execute( + async def _is_cancelled(self, job_id: uuid.UUID, session: AsyncSession) -> bool: + result = await session.execute( select(Job.cancellation_requested).where(Job.id == job_id) ) return bool(result.scalar_one_or_none()) @@ -343,34 +305,23 @@ async def _mark_cancelled( session: AsyncSession, reason: str = "cancellation_requested", ) -> None: - """Mark a job as cancelled and publish SSE event.""" await session.execute( update(Job) .where(Job.id == job_id) - .values( - status=JobStatus.CANCELLED, - worker_id=None, - updated_at=func.now(), - ) + .values(status=JobStatus.CANCELLED, worker_id=None, updated_at=func.now()) ) - await self._log_event( job_id=job_id, event=LogEvent.JOB_CANCELLED, - message=(f"Job cancelled by worker {self.worker_id}. Reason: {reason}."), + message=f"Job cancelled by worker {self.worker_id}. Reason: {reason}.", metadata={"worker_id": self.worker_id, "reason": reason}, - db=session, + session=session, ) await session.commit() - logger.info( - "job_cancelled", - job_id=str(job_id), - worker_id=self.worker_id, - reason=reason, - ) + logger.info("job_cancelled", job_id=str(job_id), reason=reason) - await _publish_sse_event( + await _publish_sse( { "job_id": str(job_id), "status": JobStatus.CANCELLED, @@ -378,19 +329,10 @@ async def _mark_cancelled( } ) - async def _handle_recurrence( - self, - job: Job, - session: AsyncSession, - ) -> None: - """ - Schedule the next run of a recurring job. - """ - + async def _handle_recurrence(self, job: Job, session: AsyncSession) -> None: if not job.interval_seconds: return - - next_run = datetime.now(UTC) + timedelta(seconds=float(job.interval_seconds)) + next_run = datetime.now(UTC) + timedelta(seconds=job.interval_seconds) new_job = Job( type=job.type, @@ -409,44 +351,43 @@ async def _handle_recurrence( await self._log_event( job_id=new_job.id, event=LogEvent.RECURRING_SCHEDULED, - message=( - f"Recurring job scheduled. Next run at {next_run.isoformat()}. " - f"Parent job: {job.id}." - ), + message=f"Next run scheduled at {next_run.isoformat()}. Parent: {job.id}.", metadata={ "parent_job_id": str(job.id), "interval_seconds": job.interval_seconds, "next_run": next_run.isoformat(), }, - db=session, + session=session, ) await session.commit() + now = datetime.now(UTC) + if next_run <= now: + await self.queue.push( + job_id=str(new_job.id), + effective_priority=float(job.priority), + scheduled_at=next_run.timestamp(), + created_at=new_job.created_at.timestamp() + if new_job.created_at + else now.timestamp(), + ) + logger.info( "recurring_job_scheduled", parent_job_id=str(job.id), new_job_id=str(new_job.id), next_run=next_run.isoformat(), - interval_seconds=job.interval_seconds, ) - async def _load_job( - self, - job_id: str, - db: AsyncSession, - ) -> Job | None: - """Load a job by string ID. Returns None if not found.""" + async def _load_job(self, job_id: str, session: AsyncSession) -> Job | None: try: parsed_id = uuid.UUID(job_id) except ValueError: logger.error("invalid_job_id", job_id=job_id) return None - result = await db.execute( - select(Job).where( - Job.id == parsed_id, - Job.deleted_at.is_(None), - ) + result = await session.execute( + select(Job).where(Job.id == parsed_id, Job.deleted_at.is_(None)) ) return result.scalar_one_or_none() @@ -456,20 +397,17 @@ async def _log_event( event: str, message: str, metadata: dict, - db: AsyncSession, + session: AsyncSession, ) -> None: - """Write a structured log entry to the job_logs table.""" log_entry = JobLog( job_id=job_id, event=event, message=message, metadata_=metadata, ) - db.add(log_entry) - await db.flush() + session.add(log_entry) + await session.flush() @staticmethod def _elapsed_ms(start: datetime) -> int: - """Return elapsed milliseconds since start.""" - delta = datetime.now(UTC) - start - return int(delta.total_seconds() * 1000) + return int((datetime.now(UTC) - start).total_seconds() * 1000) diff --git a/worker/worker.py b/worker/worker.py index 366ea4f..4ebcaa1 100644 --- a/worker/worker.py +++ b/worker/worker.py @@ -1,302 +1,95 @@ import asyncio -import os -import signal -import sys -import uuid import redis.asyncio as aioredis from app.core.config import settings -from app.core.logger import get_logger, setup_logging +from app.core.logger import get_logger from app.queues.heapq import HeapQueue -from app.queues.timing_wheel import TimingWheel from worker.processor import JobProcessor -setup_logging() logger = get_logger(__name__) -WORKER_REGISTRY_KEY = "flint:workers:{worker_id}" -WORKER_CONTROL_CHANNEL = "flint:worker:control:{worker_id}" -QUEUE_KEY = "flint:queue" -HEARTBEAT_TTL = 60 # seconds — key expires if worker dies without deregistering -HEARTBEAT_INTERVAL = 30 # seconds between heartbeat refreshes +WORKER_TTL = 60 +HEARTBEAT_INTERVAL = 15 -class FlintWorker: - def __init__(self) -> None: - self.worker_id = settings.WORKER_ID or f"worker-{str(uuid.uuid4())[:8]}" - self.running = False - self._redis: aioredis.Redis - self._queue: HeapQueue | TimingWheel | None = None - self._processor: JobProcessor - self._shutdown_event = asyncio.Event() +async def run_worker( + worker_id: str, + queue: HeapQueue, + shutdown_event: asyncio.Event, +) -> None: - async def start(self) -> None: - """ - Main entry point. Initialises resources and starts all loops. - """ - self.running = True + redis = aioredis.from_url(settings.REDIS_URL, decode_responses=True) + processor = JobProcessor(worker_id=worker_id, queue=queue) + worker_key = f"flint:workers:{worker_id}" + control_channel = f"flint:worker:control:{worker_id}" - self._redis = aioredis.from_url( - settings.REDIS_URL, - decode_responses=True, - encoding="utf-8", - ) - self._redis_pubsub = aioredis.from_url( - settings.REDIS_URL, - decode_responses=True, - encoding="utf-8", - socket_timeout=None, - socket_connect_timeout=5, - ) + # Register worker + await redis.set(worker_key, "idle", ex=WORKER_TTL) + logger.info("worker_started", worker_id=worker_id) + local_stop = asyncio.Event() - self._queue = await self._build_queue() - - await self._load_queue_from_redis() - - self._processor = JobProcessor( - worker_id=self.worker_id, - queue=self._queue, - ) - - await self._register() - - loop = asyncio.get_event_loop() - for sig in (signal.SIGTERM, signal.SIGINT): - loop.add_signal_handler(sig, self._handle_shutdown_signal) - - logger.info( - "worker_started", - worker_id=self.worker_id, - queue_strategy=self._queue.__class__.__name__, - ) - try: - await asyncio.gather( - self._poll_loop(), - self._heartbeat_loop(), - self._control_channel_loop(), - ) - finally: - await self._redis_pubsub.aclose() - await self._redis.aclose() - - async def _poll_loop(self) -> None: - """ - Core poll loop. Pops job IDs from the queue and processes them. - Sleeps for WORKER_POLL_INTERVAL when the queue is empty. - """ - while self.running: + async def _heartbeat_loop(): + while not local_stop.is_set() and not shutdown_event.is_set(): + await redis.set(worker_key, "active", ex=WORKER_TTL) try: - new_queue = await self._maybe_switch_queue() - if new_queue: - self._queue = new_queue - self._processor.queue = new_queue - await self._load_queue_from_redis() - - job_id = await self._pop_next() - if job_id: - await self._processor.process(job_id) - else: - await asyncio.sleep(settings.WORKER_POLL_INTERVAL) - - except asyncio.CancelledError: - break - except Exception as exc: - logger.error( - "worker_poll_error", - worker_id=self.worker_id, - error=str(exc), + await asyncio.wait_for( + asyncio.shield(local_stop.wait()), timeout=HEARTBEAT_INTERVAL ) - await asyncio.sleep(1.0) - - await self._deregister() - logger.info("worker_stopped", worker_id=self.worker_id) - - async def _pop_next(self) -> str | None: - """ - Pop the next job_id from the Redis sorted set. - Uses ZPOPMIN to atomically remove and return the lowest-score member. - """ - items = await self._redis.zpopmin(QUEUE_KEY, 1) - if not items: - return None - - job_id = items[0] # float(items[1]) if len(items) > 1 else float(items[0][1]), - - # Handle both tuple and flat list responses - if isinstance(items[0], (list, tuple)): - job_id = items[0][0] - # score = float(items[0][1]) - else: - job_id = items[0] - # score = float(items[1]) - - return str(job_id) - - async def _build_queue(self) -> HeapQueue | TimingWheel: - """Build the correct queue based on scheduler_strategy setting.""" - strategy = await self._get_strategy() - if strategy == "timing_wheel": - logger.info("queue_strategy", strategy="timing_wheel") - return TimingWheel() - logger.info("queue_strategy", strategy="heap") - return HeapQueue() + except TimeoutError: + pass - async def _maybe_switch_queue(self) -> HeapQueue | TimingWheel | None: - """ - Check if the scheduler_strategy setting has changed. - If it has, return a new queue instance. Otherwise return None. - """ - current_strategy = ( - self._queue.__class__.__name__.lower() - .replace("heapqueue", "heap") - .replace("timingwheel", "timing_wheel") - ) - - new_strategy = await self._get_strategy() - - if new_strategy != current_strategy: - logger.info( - "queue_strategy_switched", - from_strategy=current_strategy, - to_strategy=new_strategy, - worker_id=self.worker_id, - ) - if new_strategy == "timing_wheel": - return TimingWheel() - return HeapQueue() - return None - - async def _get_strategy(self) -> str: - """Read scheduler_strategy from Redis cache or fall back to 'heap'.""" + async def _control_loop(): + pubsub = redis.pubsub() + await pubsub.subscribe(control_channel) try: - from app.db.session import AsyncSessionLocal - from app.services.settings import get_scheduler_strategy - - async with AsyncSessionLocal() as session: - return await get_scheduler_strategy(session) - except Exception: - return "heap" - - async def _load_queue_from_redis(self) -> None: - """ - On startup or strategy switch, load all pending job IDs from - the Redis sorted set into the local in-memory queue. - This populates the heap so the worker has the full queue state. - """ - import time - - try: - # ZRANGE with WITHSCORES returns [member, score, member, score, ...] - items = await self._redis.zrange(QUEUE_KEY, 0, -1, withscores=True) - count = 0 - for i in range(0, len(items), 2): - job_id = items[i] - score = float(items[i + 1]) - now = time.time() - await self._queue.push( - job_id=job_id, - effective_priority=score, - scheduled_at=now, - created_at=now, - ) - count += 1 - logger.info( - "queue_loaded_from_redis", - worker_id=self.worker_id, - job_count=count, - ) - except Exception as exc: - logger.error( - "queue_load_error", - worker_id=self.worker_id, - error=str(exc), - ) - - async def _heartbeat_loop(self) -> None: - """ - Refresh the worker's Redis registration key every HEARTBEAT_INTERVAL. - If the worker dies, the key expires after HEARTBEAT_TTL seconds - and the API reflects the worker as gone. - """ - while self.running: - try: - await self._register() - except Exception as exc: - logger.error( - "heartbeat_error", - worker_id=self.worker_id, - error=str(exc), + async for message in pubsub.listen(): + if message["type"] != "message": + continue + command = message["data"] + logger.info( + "worker_control_command", worker_id=worker_id, command=command ) - await asyncio.sleep(HEARTBEAT_INTERVAL) + if command == "stop": + local_stop.set() + break + elif command == "restart": + local_stop.set() + break + except asyncio.CancelledError: + pass + finally: + await pubsub.unsubscribe(control_channel) + await pubsub.aclose() - async def _register(self) -> None: - """Write/refresh the worker registration key in Redis.""" - key = WORKER_REGISTRY_KEY.format(worker_id=self.worker_id) - await self._redis.set(key, "active", ex=HEARTBEAT_TTL) + heartbeat_task = asyncio.create_task( + _heartbeat_loop(), name=f"{worker_id}-heartbeat" + ) + control_task = asyncio.create_task(_control_loop(), name=f"{worker_id}-control") - async def _deregister(self) -> None: - """Remove the worker registration key on clean shutdown.""" - key = WORKER_REGISTRY_KEY.format(worker_id=self.worker_id) - await self._redis.delete(key) - logger.info("worker_deregistered", worker_id=self.worker_id) + try: + while not local_stop.is_set() and not shutdown_event.is_set(): + try: + job_id = await queue.pop() + if job_id is None: + await asyncio.sleep(settings.WORKER_POLL_INTERVAL) + continue - async def _control_channel_loop(self) -> None: - """ - Subscribe to the worker's Redis control channel. - The API publishes 'stop' or 'restart' signals here. - This allows workers to be controlled from the dashboard. - """ - channel = WORKER_CONTROL_CHANNEL.format(worker_id=self.worker_id) + await redis.set(worker_key, "busy", ex=WORKER_TTL) + logger.info("worker_picked_job", worker_id=worker_id, job_id=job_id) + await processor.process(job_id) + await redis.set(worker_key, "idle", ex=WORKER_TTL) - while self.running: - pubsub = self._redis_pubsub.pubsub() - await pubsub.subscribe(channel) - try: - async for message in pubsub.listen(): - if not self.running: - return - if message["type"] != "message": - continue - command = message.get("data", "").strip().lower() - logger.info( - "worker_control_command", - worker_id=self.worker_id, - command=command, - ) - if command == "stop": - self._handle_shutdown_signal() - return - elif command == "restart": - logger.info("worker_restarting", worker_id=self.worker_id) - self._handle_shutdown_signal() - os.execv(sys.executable, [sys.executable] + sys.argv) except asyncio.CancelledError: - return + break except Exception as exc: - logger.warning( - "control_channel_reconnecting", - worker_id=self.worker_id, - error=str(exc), - ) - await asyncio.sleep(2) - finally: - await pubsub.unsubscribe(channel) - await pubsub.aclose() - - def _handle_shutdown_signal(self) -> None: - """ - Handle SIGTERM / SIGINT. - """ - if self.running: - logger.info( - "worker_shutdown_signal", - worker_id=self.worker_id, - ) - self.running = False - self._shutdown_event.set() - - -if __name__ == "__main__": - worker = FlintWorker() - asyncio.run(worker.start()) + logger.error("worker_loop_error", worker_id=worker_id, error=str(exc)) + await redis.set(worker_key, "idle", ex=WORKER_TTL) + await asyncio.sleep(1.0) + finally: + heartbeat_task.cancel() + control_task.cancel() + await asyncio.gather(heartbeat_task, control_task, return_exceptions=True) + await redis.delete(worker_key) + await redis.aclose() + logger.info("worker_stopped", worker_id=worker_id)