fixed

videoarchiver/processor/cleanup_manager.py

@@ -2,25 +2,37 @@
 
 import logging
 import asyncio
-from enum import Enum
-from dataclasses import dataclass
-from typing import Optional, Dict, Any, List, Set
-from datetime import datetime
+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 .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 = "queue"
-    FFMPEG = "ffmpeg"
-    TASKS = "tasks"
-    RESOURCES = "resources"
+    QUEUE = auto()
+    FFMPEG = auto()
+    TASKS = auto()
+    RESOURCES = auto()
 
 class CleanupStrategy(Enum):
     """Cleanup strategies"""
-    NORMAL = "normal"
-    FORCE = "force"
-    GRACEFUL = "graceful"
+    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:
@@ -29,33 +41,64 @@ class CleanupResult:
     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"""
 
-    def __init__(self):
+    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"""
+        """
+        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"""
+        """
+        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"""
+        """
+        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({
@@ -69,15 +112,20 @@ class CleanupTracker:
             self.start_times.pop(cleanup_id)
             self.stage_results.pop(cleanup_id)
 
-    def get_cleanup_stats(self) -> Dict[str, Any]:
-        """Get cleanup statistics"""
-        return {
-            "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 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"""
@@ -116,20 +164,49 @@ class CleanupTracker:
 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,
-        ffmpeg_mgr: Optional[object] = None,
+        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"""
+        """
+        Perform normal cleanup of resources.
+        
+        Raises:
+            CleanupError: If cleanup fails
+        """
         cleanup_id = f"cleanup_{datetime.utcnow().timestamp()}"
         self.tracker.start_cleanup(cleanup_id)
         
@@ -137,35 +214,45 @@ class CleanupManager:
             logger.info("Starting normal cleanup...")
             
             # Clean up in stages
-            stages = [
-                (CleanupStage.QUEUE, self._cleanup_queue),
-                (CleanupStage.FFMPEG, self._cleanup_ffmpeg),
-                (CleanupStage.TASKS, self._cleanup_tasks)
-            ]
-
-            for stage, cleanup_func in stages:
+            for operation in self.cleanup_operations:
                 try:
                     start_time = datetime.utcnow()
-                    await cleanup_func()
+                    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, stage, duration=duration)
+                        CleanupResult(True, operation.stage, duration=duration)
                     )
-                except Exception as e:
-                    logger.error(f"Error in {stage.value} cleanup: {e}")
+                except asyncio.TimeoutError:
+                    error = f"Cleanup stage {operation.stage.value} timed out"
+                    logger.error(error)
                     self.tracker.record_stage_result(
                         cleanup_id,
-                        CleanupResult(False, stage, str(e))
+                        CleanupResult(False, operation.stage, error)
                     )
                     if self.strategy != CleanupStrategy.GRACEFUL:
-                        raise
+                        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 Exception as e:
-            logger.error(f"Error during normal cleanup: {str(e)}", exc_info=True)
+        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)
 
@@ -178,26 +265,26 @@ class CleanupManager:
             logger.info("Starting force cleanup...")
 
             # Force cleanup in stages
-            stages = [
-                (CleanupStage.QUEUE, self._force_cleanup_queue),
-                (CleanupStage.FFMPEG, self._force_cleanup_ffmpeg),
-                (CleanupStage.TASKS, self._force_cleanup_tasks)
-            ]
+            for operation in self.cleanup_operations:
+                if not operation.force_func:
+                    continue
 
-            for stage, cleanup_func in stages:
                 try:
                     start_time = datetime.utcnow()
-                    await cleanup_func()
+                    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, stage, duration=duration)
+                        CleanupResult(True, operation.stage, duration=duration)
                     )
                 except Exception as e:
-                    logger.error(f"Error in force {stage.value} cleanup: {e}")
+                    logger.error(f"Error in force {operation.stage.value} cleanup: {e}")
                     self.tracker.record_stage_result(
                         cleanup_id,
-                        CleanupResult(False, stage, str(e))
+                        CleanupResult(False, operation.stage, str(e))
                     )
 
             logger.info("Force cleanup completed")
@@ -209,6 +296,8 @@ class CleanupManager:
 
     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:
@@ -224,15 +313,22 @@ class CleanupManager:
                 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:
-            self.ffmpeg_mgr.kill_all_processes()
+            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"""
@@ -240,13 +336,31 @@ class CleanupManager:
             self._queue_task.cancel()
 
     def set_queue_task(self, task: asyncio.Task) -> None:
-        """Set the queue processing task for cleanup purposes"""
+        """
+        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"""
+        """
+        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)
+            "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
+            ]
         }