327 lines
13 KiB
Python
327 lines
13 KiB
Python
"""
|
|
Background tasks for guest cleanup and system maintenance
|
|
Fixed for event loop safety
|
|
"""
|
|
|
|
import asyncio
|
|
from datetime import datetime, timedelta, UTC
|
|
from typing import Optional, List, Dict, Any, Callable
|
|
from logger import logger
|
|
from database.manager import DatabaseManager
|
|
|
|
|
|
class BackgroundTaskManager:
|
|
"""Manages background tasks for the application using asyncio instead of threading"""
|
|
|
|
def __init__(self, database_manager: DatabaseManager):
|
|
self.database_manager = database_manager
|
|
self.running = False
|
|
self.tasks: List[asyncio.Task] = []
|
|
self.main_loop: Optional[asyncio.AbstractEventLoop] = None
|
|
|
|
async def cleanup_inactive_guests(self, inactive_hours: int = 24):
|
|
"""Clean up inactive guest sessions"""
|
|
try:
|
|
if self.database_manager.is_shutting_down:
|
|
logger.info("Skipping guest cleanup - application shutting down")
|
|
return 0
|
|
|
|
database = self.database_manager.get_database()
|
|
cleaned_count = await database.cleanup_inactive_guests(inactive_hours)
|
|
|
|
if cleaned_count > 0:
|
|
logger.info(f"🧹 Background cleanup: removed {cleaned_count} inactive guest sessions")
|
|
|
|
return cleaned_count
|
|
except Exception as e:
|
|
logger.error(f"❌ Error in guest cleanup: {e}")
|
|
return 0
|
|
|
|
async def cleanup_expired_verification_tokens(self):
|
|
"""Clean up expired email verification tokens"""
|
|
try:
|
|
if self.database_manager.is_shutting_down:
|
|
logger.info("Skipping token cleanup - application shutting down")
|
|
return 0
|
|
|
|
database = self.database_manager.get_database()
|
|
cleaned_count = await database.cleanup_expired_verification_tokens()
|
|
|
|
if cleaned_count > 0:
|
|
logger.info(f"🧹 Background cleanup: removed {cleaned_count} expired verification tokens")
|
|
|
|
return cleaned_count
|
|
except Exception as e:
|
|
logger.error(f"❌ Error in verification token cleanup: {e}")
|
|
return 0
|
|
|
|
async def update_guest_statistics(self):
|
|
"""Update guest usage statistics"""
|
|
try:
|
|
if self.database_manager.is_shutting_down:
|
|
logger.info("Skipping stats update - application shutting down")
|
|
return {}
|
|
|
|
database = self.database_manager.get_database()
|
|
stats = await database.get_guest_statistics()
|
|
|
|
# Log interesting statistics
|
|
if stats.get("total_guests", 0) > 0:
|
|
logger.info(
|
|
f"📊 Guest stats: {stats['total_guests']} total, "
|
|
f"{stats['active_last_hour']} active in last hour, "
|
|
f"{stats['converted_guests']} converted"
|
|
)
|
|
|
|
return stats
|
|
except Exception as e:
|
|
logger.error(f"❌ Error updating guest statistics: {e}")
|
|
return {}
|
|
|
|
async def cleanup_old_rate_limit_data(self, days_old: int = 7):
|
|
"""Clean up old rate limiting data"""
|
|
try:
|
|
if self.database_manager.is_shutting_down:
|
|
logger.info("Skipping rate limit cleanup - application shutting down")
|
|
return 0
|
|
|
|
# Get Redis client safely (using the event loop safe method)
|
|
from database.manager import redis_manager
|
|
|
|
redis = await redis_manager.get_client()
|
|
|
|
# Clean up rate limit keys older than specified days
|
|
cutoff_time = datetime.now(UTC) - timedelta(days=days_old)
|
|
pattern = "rate_limit:*"
|
|
|
|
cursor = 0
|
|
deleted_count = 0
|
|
|
|
while True:
|
|
cursor, keys = await redis.scan(cursor, match=pattern, count=100)
|
|
|
|
for key in keys:
|
|
# Check if key is old enough to delete
|
|
try:
|
|
ttl = await redis.ttl(key)
|
|
if ttl == -1: # No expiration set, check creation time
|
|
creation_time = await redis.hget(key, "created_at") # type: ignore
|
|
if creation_time:
|
|
creation_time = datetime.fromisoformat(creation_time).replace(tzinfo=UTC)
|
|
if creation_time < cutoff_time:
|
|
# Key is older than cutoff, delete it
|
|
await redis.delete(key)
|
|
deleted_count += 1
|
|
except Exception:
|
|
continue
|
|
|
|
if cursor == 0:
|
|
break
|
|
|
|
if deleted_count > 0:
|
|
logger.info(f"🧹 Cleaned up {deleted_count} old rate limit keys")
|
|
|
|
return deleted_count
|
|
except Exception as e:
|
|
logger.error(f"❌ Error cleaning up rate limit data: {e}")
|
|
return 0
|
|
|
|
async def cleanup_orphaned_data(self):
|
|
"""Clean up orphaned database records"""
|
|
try:
|
|
if self.database_manager.is_shutting_down:
|
|
return 0
|
|
|
|
database = self.database_manager.get_database()
|
|
|
|
# Clean up orphaned job requirements
|
|
orphaned_count = await database.cleanup_orphaned_job_requirements()
|
|
|
|
if orphaned_count > 0:
|
|
logger.info(f"🧹 Cleaned up {orphaned_count} orphaned job requirements")
|
|
|
|
return orphaned_count
|
|
except Exception as e:
|
|
logger.error(f"❌ Error cleaning up orphaned data: {e}")
|
|
return 0
|
|
|
|
async def _run_periodic_task(self, name: str, task_func: Callable, interval_seconds: int, *args, **kwargs):
|
|
"""Run a periodic task safely in the same event loop"""
|
|
logger.info(f"🔄 Starting periodic task: {name} (every {interval_seconds}s)")
|
|
|
|
while self.running:
|
|
try:
|
|
# Verify we're still in the correct event loop
|
|
current_loop = asyncio.get_running_loop()
|
|
if current_loop != self.main_loop:
|
|
logger.error(f"Task {name} detected event loop change! Stopping.")
|
|
break
|
|
|
|
# Run the task
|
|
await task_func(*args, **kwargs)
|
|
|
|
except asyncio.CancelledError:
|
|
logger.info(f"Periodic task {name} was cancelled")
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"❌ Error in periodic task {name}: {e}")
|
|
# Continue running despite errors
|
|
|
|
# Sleep with cancellation support
|
|
try:
|
|
await asyncio.sleep(interval_seconds)
|
|
except asyncio.CancelledError:
|
|
logger.info(f"Periodic task {name} cancelled during sleep")
|
|
break
|
|
|
|
async def start(self):
|
|
"""Start all background tasks in the current event loop"""
|
|
if self.running:
|
|
logger.warning("⚠️ Background task manager already running")
|
|
return
|
|
|
|
# Store the current event loop
|
|
self.main_loop = asyncio.get_running_loop()
|
|
self.running = True
|
|
|
|
# Define periodic tasks with their intervals (in seconds)
|
|
periodic_tasks = [
|
|
# (name, function, interval_seconds, *args)
|
|
("guest_cleanup", self.cleanup_inactive_guests, 6 * 3600, 48), # Every 6 hours, cleanup 48h old
|
|
("token_cleanup", self.cleanup_expired_verification_tokens, 12 * 3600), # Every 12 hours
|
|
("guest_stats", self.update_guest_statistics, 3600), # Every hour
|
|
("rate_limit_cleanup", self.cleanup_old_rate_limit_data, 24 * 3600, 7), # Daily, cleanup 7 days old
|
|
("orphaned_cleanup", self.cleanup_orphaned_data, 6 * 3600), # Every 6 hours
|
|
]
|
|
|
|
# Create asyncio tasks for each periodic task
|
|
for name, func, interval, *args in periodic_tasks:
|
|
task = asyncio.create_task(self._run_periodic_task(name, func, interval, *args), name=f"background_{name}")
|
|
self.tasks.append(task)
|
|
logger.info(f"📅 Scheduled background task: {name}")
|
|
|
|
# Run initial cleanup tasks immediately (but don't wait for them)
|
|
asyncio.create_task(self._run_initial_cleanup(), name="initial_cleanup")
|
|
|
|
logger.info("🚀 Background task manager started with asyncio tasks")
|
|
|
|
async def _run_initial_cleanup(self):
|
|
"""Run some cleanup tasks immediately on startup"""
|
|
try:
|
|
logger.info("🧹 Running initial cleanup tasks...")
|
|
|
|
# Clean up expired tokens immediately
|
|
await asyncio.sleep(5) # Give the app time to fully start
|
|
await self.cleanup_expired_verification_tokens()
|
|
|
|
# Clean up very old inactive guests (7 days old)
|
|
await self.cleanup_inactive_guests(inactive_hours=7 * 24)
|
|
|
|
# Update statistics
|
|
await self.update_guest_statistics()
|
|
|
|
logger.info("✅ Initial cleanup tasks completed")
|
|
|
|
except Exception as e:
|
|
logger.error(f"❌ Error in initial cleanup: {e}")
|
|
|
|
async def stop(self):
|
|
"""Stop all background tasks gracefully"""
|
|
logger.info("🛑 Stopping background task manager...")
|
|
|
|
self.running = False
|
|
|
|
# Cancel all running tasks
|
|
for task in self.tasks:
|
|
if not task.done():
|
|
task.cancel()
|
|
|
|
# Wait for all tasks to complete with timeout
|
|
if self.tasks:
|
|
try:
|
|
await asyncio.wait_for(asyncio.gather(*self.tasks, return_exceptions=True), timeout=30.0)
|
|
logger.info("✅ All background tasks stopped gracefully")
|
|
except asyncio.TimeoutError:
|
|
logger.warning("⚠️ Some background tasks did not stop within timeout")
|
|
|
|
self.tasks.clear()
|
|
self.main_loop = None
|
|
|
|
logger.info("🛑 Background task manager stopped")
|
|
|
|
async def get_task_status(self) -> Dict[str, Any]:
|
|
"""Get status of all background tasks"""
|
|
status = {
|
|
"running": self.running,
|
|
"main_loop_id": id(self.main_loop) if self.main_loop else None,
|
|
"current_loop_id": None,
|
|
"task_count": len(self.tasks),
|
|
"tasks": [],
|
|
}
|
|
|
|
try:
|
|
current_loop = asyncio.get_running_loop()
|
|
status["current_loop_id"] = id(current_loop)
|
|
status["loop_matches"] = (id(current_loop) == id(self.main_loop)) if self.main_loop else False
|
|
except RuntimeError:
|
|
status["current_loop_id"] = "no_running_loop"
|
|
|
|
for task in self.tasks:
|
|
task_info = {
|
|
"name": task.get_name(),
|
|
"done": task.done(),
|
|
"cancelled": task.cancelled(),
|
|
}
|
|
|
|
if task.done() and not task.cancelled():
|
|
try:
|
|
task.result() # This will raise an exception if the task failed
|
|
task_info["status"] = "completed"
|
|
except Exception as e:
|
|
task_info["status"] = "failed"
|
|
task_info["error"] = str(e)
|
|
elif task.cancelled():
|
|
task_info["status"] = "cancelled"
|
|
else:
|
|
task_info["status"] = "running"
|
|
|
|
status["tasks"].append(task_info)
|
|
|
|
return status
|
|
|
|
async def force_run_task(self, task_name: str) -> Any:
|
|
"""Manually trigger a specific background task"""
|
|
task_map = {
|
|
"guest_cleanup": self.cleanup_inactive_guests,
|
|
"token_cleanup": self.cleanup_expired_verification_tokens,
|
|
"guest_stats": self.update_guest_statistics,
|
|
"rate_limit_cleanup": self.cleanup_old_rate_limit_data,
|
|
"orphaned_cleanup": self.cleanup_orphaned_data,
|
|
}
|
|
|
|
if task_name not in task_map:
|
|
raise ValueError(f"Unknown task: {task_name}. Available: {list(task_map.keys())}")
|
|
|
|
logger.info(f"🔧 Manually running task: {task_name}")
|
|
result = await task_map[task_name]()
|
|
logger.info(f"✅ Manual task {task_name} completed")
|
|
return result
|
|
|
|
|
|
# Usage in your main application
|
|
async def setup_background_tasks(database_manager: DatabaseManager) -> BackgroundTaskManager:
|
|
"""Setup and start background tasks"""
|
|
task_manager = BackgroundTaskManager(database_manager)
|
|
await task_manager.start()
|
|
return task_manager
|
|
|
|
|
|
# For integration with your existing app startup
|
|
async def initialize_with_background_tasks(database_manager: DatabaseManager):
|
|
"""Initialize database and background tasks together"""
|
|
# Start background tasks
|
|
background_tasks = await setup_background_tasks(database_manager)
|
|
|
|
# Return both for your app to manage
|
|
return database_manager, background_tasks
|