Bryan Van Deusen
245 lines
8.0 KiB
Python
245 lines
8.0 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
|
|
|
|
|
|
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
|