"""Module for managing cleanup operations in the video processor""" import logging import asyncio from enum import Enum, auto from dataclasses import dataclass, field from typing import ( Optional, Dict, Any, List, Set, TypedDict, ClassVar, Callable, Awaitable, Tuple, ) from datetime import datetime, timedelta from ..processor.queue_handler import QueueHandler from ..ffmpeg.ffmpeg_manager import FFmpegManager from ..utils.exceptions import CleanupError logger = logging.getLogger("VideoArchiver") class CleanupStage(Enum): """Cleanup stages""" QUEUE = auto() FFMPEG = auto() TASKS = auto() RESOURCES = auto() class CleanupStrategy(Enum): """Cleanup strategies""" NORMAL = auto() FORCE = auto() GRACEFUL = auto() class CleanupStats(TypedDict): """Type definition for cleanup statistics""" total_cleanups: int active_cleanups: int success_rate: float average_duration: float stage_success_rates: Dict[str, float] @dataclass class CleanupResult: """Result of a cleanup operation""" success: bool stage: CleanupStage error: Optional[str] = None duration: float = 0.0 timestamp: str = field(default_factory=lambda: datetime.utcnow().isoformat()) @dataclass class CleanupOperation: """Represents a cleanup operation""" stage: CleanupStage func: Callable[[], Awaitable[None]] force_func: Optional[Callable[[], Awaitable[None]]] = None timeout: float = 30.0 # Default timeout in seconds class CleanupTracker: """Tracks cleanup operations""" MAX_HISTORY: ClassVar[int] = 1000 # Maximum number of cleanup operations to track def __init__(self) -> None: self.cleanup_history: List[Dict[str, Any]] = [] self.active_cleanups: Set[str] = set() self.start_times: Dict[str, datetime] = {} self.stage_results: Dict[str, List[CleanupResult]] = {} def start_cleanup(self, cleanup_id: str) -> None: """ Start tracking a cleanup operation. Args: cleanup_id: Unique identifier for the cleanup operation """ self.active_cleanups.add(cleanup_id) self.start_times[cleanup_id] = datetime.utcnow() self.stage_results[cleanup_id] = [] # Cleanup old history if needed if len(self.cleanup_history) >= self.MAX_HISTORY: self.cleanup_history = self.cleanup_history[-self.MAX_HISTORY :] def record_stage_result(self, cleanup_id: str, result: CleanupResult) -> None: """ Record result of a cleanup stage. Args: cleanup_id: Cleanup operation identifier result: Result of the cleanup stage """ if cleanup_id in self.stage_results: self.stage_results[cleanup_id].append(result) def end_cleanup(self, cleanup_id: str) -> None: """ End tracking a cleanup operation. Args: cleanup_id: Cleanup operation identifier """ if cleanup_id in self.active_cleanups: end_time = datetime.utcnow() self.cleanup_history.append( { "id": cleanup_id, "start_time": self.start_times[cleanup_id], "end_time": end_time, "duration": ( end_time - self.start_times[cleanup_id] ).total_seconds(), "results": self.stage_results[cleanup_id], } ) self.active_cleanups.remove(cleanup_id) self.start_times.pop(cleanup_id) self.stage_results.pop(cleanup_id) def get_cleanup_stats(self) -> CleanupStats: """ Get cleanup statistics. Returns: Dictionary containing cleanup statistics """ return CleanupStats( total_cleanups=len(self.cleanup_history), active_cleanups=len(self.active_cleanups), success_rate=self._calculate_success_rate(), average_duration=self._calculate_average_duration(), stage_success_rates=self._calculate_stage_success_rates(), ) def _calculate_success_rate(self) -> float: """Calculate overall cleanup success rate""" if not self.cleanup_history: return 1.0 successful = sum( 1 for cleanup in self.cleanup_history if all(result.success for result in cleanup["results"]) ) return successful / len(self.cleanup_history) def _calculate_average_duration(self) -> float: """Calculate average cleanup duration""" if not self.cleanup_history: return 0.0 total_duration = sum(cleanup["duration"] for cleanup in self.cleanup_history) return total_duration / len(self.cleanup_history) def _calculate_stage_success_rates(self) -> Dict[str, float]: """Calculate success rates by stage""" stage_attempts: Dict[str, int] = {} stage_successes: Dict[str, int] = {} for cleanup in self.cleanup_history: for result in cleanup["results"]: stage = result.stage.value stage_attempts[stage] = stage_attempts.get(stage, 0) + 1 if result.success: stage_successes[stage] = stage_successes.get(stage, 0) + 1 return { stage: stage_successes.get(stage, 0) / attempts for stage, attempts in stage_attempts.items() } class CleanupManager: """Manages cleanup operations for the video processor""" CLEANUP_TIMEOUT: ClassVar[int] = 60 # Default timeout for entire cleanup operation def __init__( self, queue_handler: QueueHandler, ffmpeg_mgr: Optional[FFmpegManager] = None, strategy: CleanupStrategy = CleanupStrategy.NORMAL, ) -> None: self.queue_handler = queue_handler self.ffmpeg_mgr = ffmpeg_mgr self.strategy = strategy self._queue_task: Optional[asyncio.Task] = None self.tracker = CleanupTracker() # Define cleanup operations self.cleanup_operations: List[CleanupOperation] = [ CleanupOperation( stage=CleanupStage.QUEUE, func=self._cleanup_queue, force_func=self._force_cleanup_queue, timeout=30.0, ), CleanupOperation( stage=CleanupStage.FFMPEG, func=self._cleanup_ffmpeg, force_func=self._force_cleanup_ffmpeg, timeout=15.0, ), CleanupOperation( stage=CleanupStage.TASKS, func=self._cleanup_tasks, force_func=self._force_cleanup_tasks, timeout=15.0, ), ] async def cleanup(self) -> None: """ Perform normal cleanup of resources. Raises: CleanupError: If cleanup fails """ cleanup_id = f"cleanup_{datetime.utcnow().timestamp()}" self.tracker.start_cleanup(cleanup_id) try: logger.info("Starting normal cleanup...") # Clean up in stages for operation in self.cleanup_operations: try: start_time = datetime.utcnow() await asyncio.wait_for(operation.func(), timeout=operation.timeout) duration = (datetime.utcnow() - start_time).total_seconds() self.tracker.record_stage_result( cleanup_id, CleanupResult(True, operation.stage, duration=duration), ) except asyncio.TimeoutError: error = f"Cleanup stage {operation.stage.value} timed out" logger.error(error) self.tracker.record_stage_result( cleanup_id, CleanupResult(False, operation.stage, error) ) if self.strategy != CleanupStrategy.GRACEFUL: raise CleanupError(error) except Exception as e: error = f"Error in {operation.stage.value} cleanup: {e}" logger.error(error) self.tracker.record_stage_result( cleanup_id, CleanupResult(False, operation.stage, str(e)) ) if self.strategy != CleanupStrategy.GRACEFUL: raise CleanupError(error) logger.info("Normal cleanup completed successfully") except CleanupError: raise except Exception as e: error = f"Unexpected error during cleanup: {str(e)}" logger.error(error, exc_info=True) raise CleanupError(error) finally: self.tracker.end_cleanup(cleanup_id) async def force_cleanup(self) -> None: """Force cleanup of resources when normal cleanup fails""" cleanup_id = f"force_cleanup_{datetime.utcnow().timestamp()}" self.tracker.start_cleanup(cleanup_id) try: logger.info("Starting force cleanup...") # Force cleanup in stages for operation in self.cleanup_operations: if not operation.force_func: continue try: start_time = datetime.utcnow() await asyncio.wait_for( operation.force_func(), timeout=operation.timeout ) duration = (datetime.utcnow() - start_time).total_seconds() self.tracker.record_stage_result( cleanup_id, CleanupResult(True, operation.stage, duration=duration), ) except Exception as e: logger.error(f"Error in force {operation.stage.value} cleanup: {e}") self.tracker.record_stage_result( cleanup_id, CleanupResult(False, operation.stage, str(e)) ) logger.info("Force cleanup completed") except Exception as e: logger.error(f"Error during force cleanup: {str(e)}", exc_info=True) finally: self.tracker.end_cleanup(cleanup_id) async def _cleanup_queue(self) -> None: """Clean up queue handler""" if not self.queue_handler: raise CleanupError("Queue handler not initialized") await self.queue_handler.cleanup() async def _cleanup_ffmpeg(self) -> None: """Clean up FFmpeg manager""" if self.ffmpeg_mgr: self.ffmpeg_mgr.kill_all_processes() async def _cleanup_tasks(self) -> None: """Clean up tasks""" if self._queue_task and not self._queue_task.done(): self._queue_task.cancel() try: await self._queue_task except asyncio.CancelledError: pass except Exception as e: raise CleanupError(f"Error cleaning up queue task: {str(e)}") async def _force_cleanup_queue(self) -> None: """Force clean up queue handler""" if not self.queue_handler: raise CleanupError("Queue handler not initialized") await self.queue_handler.force_cleanup() async def _force_cleanup_ffmpeg(self) -> None: """Force clean up FFmpeg manager""" if self.ffmpeg_mgr: try: self.ffmpeg_mgr.kill_all_processes() except Exception as e: logger.error(f"Error force cleaning FFmpeg processes: {e}") async def _force_cleanup_tasks(self) -> None: """Force clean up tasks""" if self._queue_task and not self._queue_task.done(): self._queue_task.cancel() def set_queue_task(self, task: asyncio.Task) -> None: """ Set the queue processing task for cleanup purposes. Args: task: Queue processing task to track """ self._queue_task = task def get_cleanup_stats(self) -> Dict[str, Any]: """ Get cleanup statistics. Returns: Dictionary containing cleanup statistics and status """ return { "stats": self.tracker.get_cleanup_stats(), "strategy": self.strategy.value, "active_cleanups": len(self.tracker.active_cleanups), "operations": [ { "stage": op.stage.value, "timeout": op.timeout, "has_force_cleanup": op.force_func is not None, } for op in self.cleanup_operations ], }