- Add ProgressService for centralized progress tracking and broadcasting - Integrate ProgressService with DownloadService for download progress - Integrate ProgressService with AnimeService for scan progress - Add progress-related WebSocket message models (ScanProgress, ErrorNotification, etc.) - Initialize ProgressService with WebSocket callback in application startup - Add comprehensive unit tests for ProgressService - Update infrastructure.md with ProgressService documentation - Remove completed WebSocket Real-time Updates task from instructions.md The ProgressService provides: - Real-time progress tracking for downloads, scans, and queue operations - Automatic progress percentage calculation - Progress lifecycle management (start, update, complete, fail, cancel) - WebSocket integration for instant client updates - Progress history with size limits - Thread-safe operations using asyncio locks - Support for metadata and custom messages Benefits: - Decoupled progress tracking from WebSocket broadcasting - Single reusable service across all components - Supports multiple concurrent operations efficiently - Centralized progress tracking simplifies monitoring - Instant feedback to users on long-running operations
486 lines
15 KiB
Python
486 lines
15 KiB
Python
"""Progress service for managing real-time progress updates.
|
|
|
|
This module provides a centralized service for tracking and broadcasting
|
|
real-time progress updates for downloads, scans, queue changes, and
|
|
system events. It integrates with the WebSocket service to push updates
|
|
to connected clients.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from enum import Enum
|
|
from typing import Any, Callable, Dict, Optional
|
|
|
|
import structlog
|
|
|
|
logger = structlog.get_logger(__name__)
|
|
|
|
|
|
class ProgressType(str, Enum):
|
|
"""Types of progress updates."""
|
|
|
|
DOWNLOAD = "download"
|
|
SCAN = "scan"
|
|
QUEUE = "queue"
|
|
SYSTEM = "system"
|
|
ERROR = "error"
|
|
|
|
|
|
class ProgressStatus(str, Enum):
|
|
"""Status of a progress operation."""
|
|
|
|
STARTED = "started"
|
|
IN_PROGRESS = "in_progress"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
CANCELLED = "cancelled"
|
|
|
|
|
|
@dataclass
|
|
class ProgressUpdate:
|
|
"""Represents a progress update event.
|
|
|
|
Attributes:
|
|
id: Unique identifier for this progress operation
|
|
type: Type of progress (download, scan, etc.)
|
|
status: Current status of the operation
|
|
title: Human-readable title
|
|
message: Detailed message
|
|
percent: Completion percentage (0-100)
|
|
current: Current progress value
|
|
total: Total progress value
|
|
metadata: Additional metadata
|
|
started_at: When operation started
|
|
updated_at: When last updated
|
|
"""
|
|
|
|
id: str
|
|
type: ProgressType
|
|
status: ProgressStatus
|
|
title: str
|
|
message: str = ""
|
|
percent: float = 0.0
|
|
current: int = 0
|
|
total: int = 0
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
|
started_at: datetime = field(default_factory=datetime.utcnow)
|
|
updated_at: datetime = field(default_factory=datetime.utcnow)
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert progress update to dictionary."""
|
|
return {
|
|
"id": self.id,
|
|
"type": self.type.value,
|
|
"status": self.status.value,
|
|
"title": self.title,
|
|
"message": self.message,
|
|
"percent": round(self.percent, 2),
|
|
"current": self.current,
|
|
"total": self.total,
|
|
"metadata": self.metadata,
|
|
"started_at": self.started_at.isoformat(),
|
|
"updated_at": self.updated_at.isoformat(),
|
|
}
|
|
|
|
|
|
class ProgressServiceError(Exception):
|
|
"""Service-level exception for progress operations."""
|
|
|
|
|
|
class ProgressService:
|
|
"""Manages real-time progress updates and broadcasting.
|
|
|
|
Features:
|
|
- Track multiple concurrent progress operations
|
|
- Calculate progress percentages and rates
|
|
- Broadcast updates via WebSocket
|
|
- Manage progress lifecycle (start, update, complete, fail)
|
|
- Support for different progress types (download, scan, queue)
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Initialize the progress service."""
|
|
# Active progress operations: id -> ProgressUpdate
|
|
self._active_progress: Dict[str, ProgressUpdate] = {}
|
|
|
|
# Completed progress history (limited size)
|
|
self._history: Dict[str, ProgressUpdate] = {}
|
|
self._max_history_size = 50
|
|
|
|
# WebSocket broadcast callback
|
|
self._broadcast_callback: Optional[Callable] = None
|
|
|
|
# Lock for thread-safe operations
|
|
self._lock = asyncio.Lock()
|
|
|
|
logger.info("ProgressService initialized")
|
|
|
|
def set_broadcast_callback(self, callback: Callable) -> None:
|
|
"""Set callback for broadcasting progress updates via WebSocket.
|
|
|
|
Args:
|
|
callback: Async function to call for broadcasting updates
|
|
"""
|
|
self._broadcast_callback = callback
|
|
logger.debug("Progress broadcast callback registered")
|
|
|
|
async def _broadcast(self, update: ProgressUpdate, room: str) -> None:
|
|
"""Broadcast progress update to WebSocket clients.
|
|
|
|
Args:
|
|
update: Progress update to broadcast
|
|
room: WebSocket room to broadcast to
|
|
"""
|
|
if self._broadcast_callback:
|
|
try:
|
|
await self._broadcast_callback(
|
|
message_type=f"{update.type.value}_progress",
|
|
data=update.to_dict(),
|
|
room=room,
|
|
)
|
|
except Exception as e:
|
|
logger.error(
|
|
"Failed to broadcast progress update",
|
|
error=str(e),
|
|
progress_id=update.id,
|
|
)
|
|
|
|
async def start_progress(
|
|
self,
|
|
progress_id: str,
|
|
progress_type: ProgressType,
|
|
title: str,
|
|
total: int = 0,
|
|
message: str = "",
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
) -> ProgressUpdate:
|
|
"""Start a new progress operation.
|
|
|
|
Args:
|
|
progress_id: Unique identifier for this progress
|
|
progress_type: Type of progress operation
|
|
title: Human-readable title
|
|
total: Total items/bytes to process
|
|
message: Initial message
|
|
metadata: Additional metadata
|
|
|
|
Returns:
|
|
Created progress update object
|
|
|
|
Raises:
|
|
ProgressServiceError: If progress already exists
|
|
"""
|
|
async with self._lock:
|
|
if progress_id in self._active_progress:
|
|
raise ProgressServiceError(
|
|
f"Progress with id '{progress_id}' already exists"
|
|
)
|
|
|
|
update = ProgressUpdate(
|
|
id=progress_id,
|
|
type=progress_type,
|
|
status=ProgressStatus.STARTED,
|
|
title=title,
|
|
message=message,
|
|
total=total,
|
|
metadata=metadata or {},
|
|
)
|
|
|
|
self._active_progress[progress_id] = update
|
|
|
|
logger.info(
|
|
"Progress started",
|
|
progress_id=progress_id,
|
|
type=progress_type.value,
|
|
title=title,
|
|
)
|
|
|
|
# Broadcast to appropriate room
|
|
room = f"{progress_type.value}_progress"
|
|
await self._broadcast(update, room)
|
|
|
|
return update
|
|
|
|
async def update_progress(
|
|
self,
|
|
progress_id: str,
|
|
current: Optional[int] = None,
|
|
total: Optional[int] = None,
|
|
message: Optional[str] = None,
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
force_broadcast: bool = False,
|
|
) -> ProgressUpdate:
|
|
"""Update an existing progress operation.
|
|
|
|
Args:
|
|
progress_id: Progress identifier
|
|
current: Current progress value
|
|
total: Updated total value
|
|
message: Updated message
|
|
metadata: Additional metadata to merge
|
|
force_broadcast: Force broadcasting even for small changes
|
|
|
|
Returns:
|
|
Updated progress object
|
|
|
|
Raises:
|
|
ProgressServiceError: If progress not found
|
|
"""
|
|
async with self._lock:
|
|
if progress_id not in self._active_progress:
|
|
raise ProgressServiceError(
|
|
f"Progress with id '{progress_id}' not found"
|
|
)
|
|
|
|
update = self._active_progress[progress_id]
|
|
old_percent = update.percent
|
|
|
|
# Update fields
|
|
if current is not None:
|
|
update.current = current
|
|
if total is not None:
|
|
update.total = total
|
|
if message is not None:
|
|
update.message = message
|
|
if metadata:
|
|
update.metadata.update(metadata)
|
|
|
|
# Calculate percentage
|
|
if update.total > 0:
|
|
update.percent = (update.current / update.total) * 100
|
|
else:
|
|
update.percent = 0.0
|
|
|
|
update.status = ProgressStatus.IN_PROGRESS
|
|
update.updated_at = datetime.utcnow()
|
|
|
|
# Only broadcast if significant change or forced
|
|
percent_change = abs(update.percent - old_percent)
|
|
should_broadcast = force_broadcast or percent_change >= 1.0
|
|
|
|
if should_broadcast:
|
|
room = f"{update.type.value}_progress"
|
|
await self._broadcast(update, room)
|
|
|
|
return update
|
|
|
|
async def complete_progress(
|
|
self,
|
|
progress_id: str,
|
|
message: str = "Completed successfully",
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
) -> ProgressUpdate:
|
|
"""Mark a progress operation as completed.
|
|
|
|
Args:
|
|
progress_id: Progress identifier
|
|
message: Completion message
|
|
metadata: Additional metadata
|
|
|
|
Returns:
|
|
Completed progress object
|
|
|
|
Raises:
|
|
ProgressServiceError: If progress not found
|
|
"""
|
|
async with self._lock:
|
|
if progress_id not in self._active_progress:
|
|
raise ProgressServiceError(
|
|
f"Progress with id '{progress_id}' not found"
|
|
)
|
|
|
|
update = self._active_progress[progress_id]
|
|
update.status = ProgressStatus.COMPLETED
|
|
update.message = message
|
|
update.percent = 100.0
|
|
update.current = update.total
|
|
update.updated_at = datetime.utcnow()
|
|
|
|
if metadata:
|
|
update.metadata.update(metadata)
|
|
|
|
# Move to history
|
|
del self._active_progress[progress_id]
|
|
self._add_to_history(update)
|
|
|
|
logger.info(
|
|
"Progress completed",
|
|
progress_id=progress_id,
|
|
type=update.type.value,
|
|
)
|
|
|
|
# Broadcast completion
|
|
room = f"{update.type.value}_progress"
|
|
await self._broadcast(update, room)
|
|
|
|
return update
|
|
|
|
async def fail_progress(
|
|
self,
|
|
progress_id: str,
|
|
error_message: str,
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
) -> ProgressUpdate:
|
|
"""Mark a progress operation as failed.
|
|
|
|
Args:
|
|
progress_id: Progress identifier
|
|
error_message: Error description
|
|
metadata: Additional error metadata
|
|
|
|
Returns:
|
|
Failed progress object
|
|
|
|
Raises:
|
|
ProgressServiceError: If progress not found
|
|
"""
|
|
async with self._lock:
|
|
if progress_id not in self._active_progress:
|
|
raise ProgressServiceError(
|
|
f"Progress with id '{progress_id}' not found"
|
|
)
|
|
|
|
update = self._active_progress[progress_id]
|
|
update.status = ProgressStatus.FAILED
|
|
update.message = error_message
|
|
update.updated_at = datetime.utcnow()
|
|
|
|
if metadata:
|
|
update.metadata.update(metadata)
|
|
|
|
# Move to history
|
|
del self._active_progress[progress_id]
|
|
self._add_to_history(update)
|
|
|
|
logger.error(
|
|
"Progress failed",
|
|
progress_id=progress_id,
|
|
type=update.type.value,
|
|
error=error_message,
|
|
)
|
|
|
|
# Broadcast failure
|
|
room = f"{update.type.value}_progress"
|
|
await self._broadcast(update, room)
|
|
|
|
return update
|
|
|
|
async def cancel_progress(
|
|
self,
|
|
progress_id: str,
|
|
message: str = "Cancelled by user",
|
|
) -> ProgressUpdate:
|
|
"""Cancel a progress operation.
|
|
|
|
Args:
|
|
progress_id: Progress identifier
|
|
message: Cancellation message
|
|
|
|
Returns:
|
|
Cancelled progress object
|
|
|
|
Raises:
|
|
ProgressServiceError: If progress not found
|
|
"""
|
|
async with self._lock:
|
|
if progress_id not in self._active_progress:
|
|
raise ProgressServiceError(
|
|
f"Progress with id '{progress_id}' not found"
|
|
)
|
|
|
|
update = self._active_progress[progress_id]
|
|
update.status = ProgressStatus.CANCELLED
|
|
update.message = message
|
|
update.updated_at = datetime.utcnow()
|
|
|
|
# Move to history
|
|
del self._active_progress[progress_id]
|
|
self._add_to_history(update)
|
|
|
|
logger.info(
|
|
"Progress cancelled",
|
|
progress_id=progress_id,
|
|
type=update.type.value,
|
|
)
|
|
|
|
# Broadcast cancellation
|
|
room = f"{update.type.value}_progress"
|
|
await self._broadcast(update, room)
|
|
|
|
return update
|
|
|
|
def _add_to_history(self, update: ProgressUpdate) -> None:
|
|
"""Add completed progress to history with size limit."""
|
|
self._history[update.id] = update
|
|
|
|
# Maintain history size limit
|
|
if len(self._history) > self._max_history_size:
|
|
# Remove oldest entries
|
|
oldest_keys = sorted(
|
|
self._history.keys(),
|
|
key=lambda k: self._history[k].updated_at,
|
|
)[: len(self._history) - self._max_history_size]
|
|
|
|
for key in oldest_keys:
|
|
del self._history[key]
|
|
|
|
async def get_progress(self, progress_id: str) -> Optional[ProgressUpdate]:
|
|
"""Get current progress state.
|
|
|
|
Args:
|
|
progress_id: Progress identifier
|
|
|
|
Returns:
|
|
Progress update object or None if not found
|
|
"""
|
|
async with self._lock:
|
|
if progress_id in self._active_progress:
|
|
return self._active_progress[progress_id]
|
|
if progress_id in self._history:
|
|
return self._history[progress_id]
|
|
return None
|
|
|
|
async def get_all_active_progress(
|
|
self, progress_type: Optional[ProgressType] = None
|
|
) -> Dict[str, ProgressUpdate]:
|
|
"""Get all active progress operations.
|
|
|
|
Args:
|
|
progress_type: Optional filter by progress type
|
|
|
|
Returns:
|
|
Dictionary of progress_id -> ProgressUpdate
|
|
"""
|
|
async with self._lock:
|
|
if progress_type:
|
|
return {
|
|
pid: update
|
|
for pid, update in self._active_progress.items()
|
|
if update.type == progress_type
|
|
}
|
|
return self._active_progress.copy()
|
|
|
|
async def clear_history(self) -> None:
|
|
"""Clear progress history."""
|
|
async with self._lock:
|
|
self._history.clear()
|
|
logger.info("Progress history cleared")
|
|
|
|
|
|
# Global singleton instance
|
|
_progress_service: Optional[ProgressService] = None
|
|
|
|
|
|
def get_progress_service() -> ProgressService:
|
|
"""Get or create the global progress service instance.
|
|
|
|
Returns:
|
|
Global ProgressService instance
|
|
"""
|
|
global _progress_service
|
|
if _progress_service is None:
|
|
_progress_service = ProgressService()
|
|
return _progress_service
|