Bryan Van Deusen
411 lines
15 KiB
Python
411 lines
15 KiB
Python
"""System maintenance Celery tasks."""
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import select, update
|
|
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
|
|
|
from app.config import get_settings
|
|
from app.tasks.celery_app import celery_app
|
|
from app.models.setting import Setting, STORAGE_STATS_SETTING
|
|
from app.models.download import Download, DownloadStatus
|
|
from app.models.base import utcnow
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
settings = get_settings()
|
|
|
|
# Maximum time a job can be in "running" state before considered orphaned
|
|
# This should be longer than the task_time_limit (1 hour) to account for legitimate long runs
|
|
STALE_RUNNING_THRESHOLD_MINUTES = 90
|
|
|
|
# Maximum time a job can be in "queued" state before considered lost
|
|
# If a queued job hasn't been picked up in this time, the Celery task was probably lost
|
|
# (e.g., due to Redis DB mismatch or Redis restart)
|
|
STALE_QUEUED_THRESHOLD_MINUTES = 30
|
|
|
|
|
|
def get_async_session():
|
|
"""Get async session factory for Celery tasks."""
|
|
engine = create_async_engine(settings.async_database_url, echo=False)
|
|
session_factory = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
return session_factory, engine
|
|
|
|
|
|
async def _cleanup_engine(engine):
|
|
"""Properly dispose of the async engine."""
|
|
await engine.dispose()
|
|
|
|
|
|
def format_size(size_bytes: int) -> str:
|
|
"""Format bytes as human-readable size."""
|
|
for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
|
|
if size_bytes < 1024:
|
|
return f"{size_bytes:.1f} {unit}"
|
|
size_bytes /= 1024
|
|
return f"{size_bytes:.1f} PB"
|
|
|
|
|
|
def calculate_storage_stats() -> dict:
|
|
"""Calculate storage statistics for the downloads directory.
|
|
|
|
This runs synchronously since it's a filesystem operation.
|
|
"""
|
|
downloads_path = Path(settings.download_path)
|
|
|
|
if not downloads_path.exists():
|
|
return {
|
|
"total_size": 0,
|
|
"total_size_formatted": "0 B",
|
|
"file_count": 0,
|
|
"calculated_at": datetime.utcnow().isoformat(),
|
|
}
|
|
|
|
total_size = 0
|
|
file_count = 0
|
|
|
|
try:
|
|
for root, dirs, files in os.walk(downloads_path):
|
|
for f in files:
|
|
fp = os.path.join(root, f)
|
|
try:
|
|
total_size += os.path.getsize(fp)
|
|
file_count += 1
|
|
except OSError:
|
|
pass
|
|
except Exception as e:
|
|
logger.error(f"Error calculating storage stats: {e}")
|
|
return {
|
|
"total_size": 0,
|
|
"total_size_formatted": "Error",
|
|
"file_count": 0,
|
|
"error": str(e),
|
|
"calculated_at": datetime.utcnow().isoformat(),
|
|
}
|
|
|
|
return {
|
|
"total_size": total_size,
|
|
"total_size_formatted": format_size(total_size),
|
|
"file_count": file_count,
|
|
"calculated_at": datetime.utcnow().isoformat(),
|
|
}
|
|
|
|
|
|
async def _update_storage_stats_async() -> dict:
|
|
"""Async implementation to calculate and store storage stats."""
|
|
session_factory, engine = get_async_session()
|
|
|
|
try:
|
|
# Calculate stats (synchronous filesystem operation)
|
|
stats = calculate_storage_stats()
|
|
logger.info(f"Calculated storage stats: {stats['total_size_formatted']}, {stats['file_count']} files")
|
|
|
|
# Store in database
|
|
async with session_factory() as session:
|
|
# Check if setting exists
|
|
result = await session.execute(
|
|
select(Setting).where(Setting.key == STORAGE_STATS_SETTING)
|
|
)
|
|
setting = result.scalar_one_or_none()
|
|
|
|
if setting:
|
|
setting.value = stats
|
|
else:
|
|
setting = Setting(key=STORAGE_STATS_SETTING, value=stats)
|
|
session.add(setting)
|
|
|
|
await session.commit()
|
|
|
|
return stats
|
|
|
|
finally:
|
|
await _cleanup_engine(engine)
|
|
|
|
|
|
@celery_app.task(name="tasks.update_storage_stats")
|
|
def update_storage_stats() -> dict:
|
|
"""Celery task to calculate and cache storage stats.
|
|
|
|
This task should be scheduled to run periodically (e.g., every hour).
|
|
"""
|
|
logger.info("Starting storage stats calculation...")
|
|
result = asyncio.run(_update_storage_stats_async())
|
|
logger.info(f"Storage stats updated: {result}")
|
|
return result
|
|
|
|
|
|
async def _reset_orphaned_running_jobs_async(threshold_minutes: int = None) -> dict:
|
|
"""Reset jobs stuck in 'running' status back to 'queued'.
|
|
|
|
Jobs can get stuck in 'running' status if the worker crashes or restarts
|
|
while processing them. This function detects and resets these orphaned jobs.
|
|
|
|
Args:
|
|
threshold_minutes: Jobs running longer than this are considered orphaned.
|
|
If None, uses STALE_RUNNING_THRESHOLD_MINUTES.
|
|
|
|
Returns:
|
|
Dict with count of reset jobs
|
|
"""
|
|
if threshold_minutes is None:
|
|
threshold_minutes = STALE_RUNNING_THRESHOLD_MINUTES
|
|
|
|
session_factory, engine = get_async_session()
|
|
|
|
try:
|
|
async with session_factory() as session:
|
|
now = utcnow()
|
|
cutoff = now - timedelta(minutes=threshold_minutes)
|
|
|
|
# Find running jobs that started before the cutoff
|
|
# OR running jobs with no started_at (shouldn't happen but handle it)
|
|
query = select(Download).where(
|
|
Download.status == DownloadStatus.RUNNING,
|
|
(Download.started_at < cutoff) | (Download.started_at == None)
|
|
)
|
|
|
|
result = await session.execute(query)
|
|
orphaned_jobs = result.scalars().all()
|
|
|
|
reset_count = 0
|
|
for job in orphaned_jobs:
|
|
job.status = DownloadStatus.QUEUED
|
|
job.started_at = None
|
|
job.error_message = f"Reset from orphaned running state (was running since {job.started_at})"
|
|
reset_count += 1
|
|
logger.info(f"Reset orphaned job {job.id} (URL: {job.url[:50]}...) to QUEUED")
|
|
|
|
if reset_count > 0:
|
|
await session.commit()
|
|
logger.info(f"Reset {reset_count} orphaned running jobs to QUEUED")
|
|
else:
|
|
logger.debug("No orphaned running jobs found")
|
|
|
|
return {"reset_count": reset_count, "threshold_minutes": threshold_minutes}
|
|
|
|
finally:
|
|
await _cleanup_engine(engine)
|
|
|
|
|
|
async def _reset_all_running_jobs_async() -> dict:
|
|
"""Reset ALL jobs in 'running' status back to 'queued'.
|
|
|
|
This is useful on worker startup to clean up any jobs that were
|
|
interrupted by a previous shutdown/crash.
|
|
|
|
Returns:
|
|
Dict with count of reset jobs
|
|
"""
|
|
session_factory, engine = get_async_session()
|
|
|
|
try:
|
|
async with session_factory() as session:
|
|
# Find all running jobs
|
|
query = select(Download).where(Download.status == DownloadStatus.RUNNING)
|
|
result = await session.execute(query)
|
|
running_jobs = result.scalars().all()
|
|
|
|
reset_count = 0
|
|
for job in running_jobs:
|
|
job.status = DownloadStatus.QUEUED
|
|
job.started_at = None
|
|
job.error_message = "Reset on worker startup"
|
|
reset_count += 1
|
|
logger.info(f"Reset running job {job.id} to QUEUED on startup")
|
|
|
|
if reset_count > 0:
|
|
await session.commit()
|
|
logger.info(f"Reset {reset_count} running jobs to QUEUED on worker startup")
|
|
|
|
return {"reset_count": reset_count}
|
|
|
|
finally:
|
|
await _cleanup_engine(engine)
|
|
|
|
|
|
@celery_app.task(name="tasks.reset_orphaned_jobs")
|
|
def reset_orphaned_jobs(threshold_minutes: int = None) -> dict:
|
|
"""Celery task to reset orphaned running jobs.
|
|
|
|
Jobs that have been in 'running' status for longer than the threshold
|
|
are considered orphaned (worker probably crashed) and reset to 'queued'.
|
|
"""
|
|
logger.info("Checking for orphaned running jobs...")
|
|
result = asyncio.run(_reset_orphaned_running_jobs_async(threshold_minutes))
|
|
return result
|
|
|
|
|
|
def reset_all_running_jobs_sync() -> dict:
|
|
"""Synchronous function to reset all running jobs.
|
|
|
|
Called on worker startup before Celery is fully initialized.
|
|
"""
|
|
logger.info("Resetting all running jobs on worker startup...")
|
|
result = asyncio.run(_reset_all_running_jobs_async())
|
|
return result
|
|
|
|
|
|
async def _requeue_stale_queued_jobs_async(threshold_minutes: int = None) -> dict:
|
|
"""Re-queue Celery tasks for downloads stuck in 'queued' status.
|
|
|
|
Downloads can get stuck in 'queued' status if the Celery task was lost
|
|
(e.g., sent to wrong Redis DB, Redis restart, etc.). This function
|
|
detects these stale queued jobs and re-sends the Celery tasks.
|
|
|
|
Only one download per source is re-queued to prevent duplicate work.
|
|
|
|
Args:
|
|
threshold_minutes: Jobs queued longer than this are considered stale.
|
|
If None, uses STALE_QUEUED_THRESHOLD_MINUTES.
|
|
|
|
Returns:
|
|
Dict with count of re-queued jobs
|
|
"""
|
|
if threshold_minutes is None:
|
|
threshold_minutes = STALE_QUEUED_THRESHOLD_MINUTES
|
|
|
|
# Import here to avoid circular imports
|
|
from app.tasks.downloads import download_source
|
|
|
|
session_factory, engine = get_async_session()
|
|
|
|
try:
|
|
async with session_factory() as session:
|
|
now = utcnow()
|
|
cutoff = now - timedelta(minutes=threshold_minutes)
|
|
|
|
# Find queued jobs created before the cutoff, ordered by creation time
|
|
# so we process older jobs first
|
|
query = select(Download).where(
|
|
Download.status == DownloadStatus.QUEUED,
|
|
Download.created_at < cutoff
|
|
).order_by(Download.created_at.asc())
|
|
|
|
result = await session.execute(query)
|
|
stale_jobs = result.scalars().all()
|
|
|
|
# Track which sources we've already queued to avoid duplicates
|
|
queued_source_ids = set()
|
|
requeued_count = 0
|
|
skipped_duplicates = 0
|
|
|
|
for job in stale_jobs:
|
|
# Skip if we've already queued a job for this source
|
|
if job.source_id in queued_source_ids:
|
|
# Mark duplicate as skipped so it won't be checked again
|
|
job.status = DownloadStatus.SKIPPED
|
|
job.error_message = "Duplicate queued job - another job for this source was re-queued"
|
|
skipped_duplicates += 1
|
|
logger.info(f"Marked duplicate job {job.id} as SKIPPED (source {job.source_id} already re-queued)")
|
|
continue
|
|
|
|
try:
|
|
# Re-send the Celery task
|
|
download_source.delay(job.source_id, download_id=job.id)
|
|
queued_source_ids.add(job.source_id)
|
|
requeued_count += 1
|
|
logger.info(f"Re-queued stale job {job.id} (source_id: {job.source_id}, queued since: {job.created_at})")
|
|
except Exception as e:
|
|
logger.error(f"Failed to re-queue job {job.id}: {e}")
|
|
|
|
# Commit the status changes for skipped duplicates
|
|
if skipped_duplicates > 0:
|
|
await session.commit()
|
|
|
|
if requeued_count > 0:
|
|
logger.info(f"Re-queued {requeued_count} stale queued jobs (marked {skipped_duplicates} duplicates as SKIPPED)")
|
|
else:
|
|
logger.debug("No stale queued jobs found")
|
|
|
|
return {"requeued_count": requeued_count, "skipped_duplicates": skipped_duplicates, "threshold_minutes": threshold_minutes}
|
|
|
|
finally:
|
|
await _cleanup_engine(engine)
|
|
|
|
|
|
async def _requeue_all_queued_jobs_async() -> dict:
|
|
"""Re-queue Celery tasks for ALL downloads in 'queued' status.
|
|
|
|
This is useful on scheduler startup to recover from Redis DB mismatch
|
|
or other issues that caused Celery tasks to be lost.
|
|
|
|
Only one download per source is re-queued to prevent duplicate work.
|
|
|
|
Returns:
|
|
Dict with count of re-queued jobs
|
|
"""
|
|
# Import here to avoid circular imports
|
|
from app.tasks.downloads import download_source
|
|
|
|
session_factory, engine = get_async_session()
|
|
|
|
try:
|
|
async with session_factory() as session:
|
|
# Find all queued jobs, ordered by creation time (oldest first)
|
|
query = select(Download).where(
|
|
Download.status == DownloadStatus.QUEUED
|
|
).order_by(Download.created_at.asc())
|
|
result = await session.execute(query)
|
|
queued_jobs = result.scalars().all()
|
|
|
|
# Track which sources we've already queued to avoid duplicates
|
|
queued_source_ids = set()
|
|
requeued_count = 0
|
|
skipped_duplicates = 0
|
|
|
|
for job in queued_jobs:
|
|
# Skip if we've already queued a job for this source
|
|
if job.source_id in queued_source_ids:
|
|
# Mark duplicate as skipped so it won't be checked again
|
|
job.status = DownloadStatus.SKIPPED
|
|
job.error_message = "Duplicate queued job - another job for this source was re-queued"
|
|
skipped_duplicates += 1
|
|
logger.info(f"Marked duplicate job {job.id} as SKIPPED (source {job.source_id} already re-queued)")
|
|
continue
|
|
|
|
try:
|
|
download_source.delay(job.source_id, download_id=job.id)
|
|
queued_source_ids.add(job.source_id)
|
|
requeued_count += 1
|
|
logger.info(f"Re-queued job {job.id} (source_id: {job.source_id}) on startup")
|
|
except Exception as e:
|
|
logger.error(f"Failed to re-queue job {job.id}: {e}")
|
|
|
|
# Commit the status changes for skipped duplicates
|
|
if skipped_duplicates > 0:
|
|
await session.commit()
|
|
|
|
if requeued_count > 0:
|
|
logger.info(f"Re-queued {requeued_count} queued jobs on startup (marked {skipped_duplicates} duplicates as SKIPPED)")
|
|
|
|
return {"requeued_count": requeued_count, "skipped_duplicates": skipped_duplicates}
|
|
|
|
finally:
|
|
await _cleanup_engine(engine)
|
|
|
|
|
|
@celery_app.task(name="tasks.requeue_stale_jobs")
|
|
def requeue_stale_jobs(threshold_minutes: int = None) -> dict:
|
|
"""Celery task to re-queue stale queued jobs.
|
|
|
|
Jobs that have been in 'queued' status for longer than the threshold
|
|
are considered lost (Celery task missing) and re-queued.
|
|
"""
|
|
logger.info("Checking for stale queued jobs...")
|
|
result = asyncio.run(_requeue_stale_queued_jobs_async(threshold_minutes))
|
|
return result
|
|
|
|
|
|
def requeue_all_queued_jobs_sync() -> dict:
|
|
"""Synchronous function to re-queue all queued jobs.
|
|
|
|
Called on scheduler startup to recover lost Celery tasks.
|
|
"""
|
|
logger.info("Re-queuing all queued jobs on scheduler startup...")
|
|
result = asyncio.run(_requeue_all_queued_jobs_async())
|
|
return result
|