"""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 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