549 lines
20 KiB
Python
549 lines
20 KiB
Python
"""
|
|
Performance & Optimization Module for AniWorld App
|
|
|
|
This module provides download speed limiting, parallel download support,
|
|
caching mechanisms, memory usage monitoring, and download resumption.
|
|
"""
|
|
|
|
import os
|
|
import threading
|
|
import time
|
|
import logging
|
|
import queue
|
|
import hashlib
|
|
from datetime import datetime, timedelta
|
|
from typing import Dict, List, Optional, Any, Callable
|
|
from dataclasses import dataclass, field
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
import json
|
|
import sqlite3
|
|
from contextlib import contextmanager
|
|
import gc
|
|
import psutil
|
|
import requests
|
|
|
|
|
|
@dataclass
|
|
class DownloadTask:
|
|
"""Represents a download task with all necessary information."""
|
|
task_id: str
|
|
serie_name: str
|
|
season: int
|
|
episode: int
|
|
key: str
|
|
language: str
|
|
output_path: str
|
|
temp_path: str
|
|
priority: int = 0 # Higher number = higher priority
|
|
retry_count: int = 0
|
|
max_retries: int = 3
|
|
created_at: datetime = field(default_factory=datetime.now)
|
|
started_at: Optional[datetime] = None
|
|
completed_at: Optional[datetime] = None
|
|
status: str = 'pending' # pending, downloading, completed, failed, paused
|
|
progress: Dict[str, Any] = field(default_factory=dict)
|
|
error_message: Optional[str] = None
|
|
|
|
|
|
class SpeedLimiter:
|
|
"""Control download speeds to prevent bandwidth saturation."""
|
|
|
|
def __init__(self, max_speed_mbps: float = 0): # 0 = unlimited
|
|
self.max_speed_mbps = max_speed_mbps
|
|
self.max_bytes_per_second = max_speed_mbps * 1024 * 1024 if max_speed_mbps > 0 else 0
|
|
self.download_start_time = None
|
|
self.bytes_downloaded = 0
|
|
self.lock = threading.Lock()
|
|
self.logger = logging.getLogger(__name__)
|
|
|
|
def set_speed_limit(self, max_speed_mbps: float):
|
|
"""Set maximum download speed in MB/s."""
|
|
with self.lock:
|
|
self.max_speed_mbps = max_speed_mbps
|
|
self.max_bytes_per_second = max_speed_mbps * 1024 * 1024 if max_speed_mbps > 0 else 0
|
|
self.logger.info(f"Speed limit set to {max_speed_mbps} MB/s")
|
|
|
|
def start_download(self):
|
|
"""Mark the start of a new download session."""
|
|
with self.lock:
|
|
self.download_start_time = time.time()
|
|
self.bytes_downloaded = 0
|
|
|
|
def update_progress(self, bytes_downloaded: int):
|
|
"""Update download progress and apply speed limiting if needed."""
|
|
if self.max_bytes_per_second <= 0: # No limit
|
|
return
|
|
|
|
with self.lock:
|
|
self.bytes_downloaded += bytes_downloaded
|
|
|
|
if self.download_start_time:
|
|
elapsed_time = time.time() - self.download_start_time
|
|
if elapsed_time > 0:
|
|
current_speed = self.bytes_downloaded / elapsed_time
|
|
|
|
if current_speed > self.max_bytes_per_second:
|
|
# Calculate required delay
|
|
target_time = self.bytes_downloaded / self.max_bytes_per_second
|
|
delay = target_time - elapsed_time
|
|
|
|
if delay > 0:
|
|
self.logger.debug(f"Speed limiting: sleeping for {delay:.2f}s")
|
|
time.sleep(delay)
|
|
|
|
def get_current_speed(self) -> float:
|
|
"""Get current download speed in MB/s."""
|
|
with self.lock:
|
|
if self.download_start_time:
|
|
elapsed_time = time.time() - self.download_start_time
|
|
if elapsed_time > 0:
|
|
speed_bps = self.bytes_downloaded / elapsed_time
|
|
return speed_bps / (1024 * 1024) # Convert to MB/s
|
|
return 0.0
|
|
|
|
class MemoryMonitor:
|
|
"""Monitor and optimize memory usage."""
|
|
|
|
def __init__(self, warning_threshold_mb: int = 1024, critical_threshold_mb: int = 2048):
|
|
self.warning_threshold = warning_threshold_mb * 1024 * 1024
|
|
self.critical_threshold = critical_threshold_mb * 1024 * 1024
|
|
self.logger = logging.getLogger(__name__)
|
|
self.monitoring = False
|
|
self.monitor_thread = None
|
|
|
|
def start_monitoring(self, check_interval: int = 30):
|
|
"""Start continuous memory monitoring."""
|
|
if self.monitoring:
|
|
return
|
|
|
|
self.monitoring = True
|
|
self.monitor_thread = threading.Thread(
|
|
target=self._monitoring_loop,
|
|
args=(check_interval,),
|
|
daemon=True
|
|
)
|
|
self.monitor_thread.start()
|
|
self.logger.info("Memory monitoring started")
|
|
|
|
def stop_monitoring(self):
|
|
"""Stop memory monitoring."""
|
|
self.monitoring = False
|
|
if self.monitor_thread:
|
|
self.monitor_thread.join(timeout=5)
|
|
self.logger.info("Memory monitoring stopped")
|
|
|
|
def _monitoring_loop(self, check_interval: int):
|
|
"""Main monitoring loop."""
|
|
while self.monitoring:
|
|
try:
|
|
self.check_memory_usage()
|
|
time.sleep(check_interval)
|
|
except Exception as e:
|
|
self.logger.error(f"Error in memory monitoring: {e}")
|
|
time.sleep(check_interval)
|
|
|
|
def check_memory_usage(self):
|
|
"""Check current memory usage and take action if needed."""
|
|
try:
|
|
process = psutil.Process()
|
|
memory_info = process.memory_info()
|
|
memory_usage = memory_info.rss
|
|
|
|
if memory_usage > self.critical_threshold:
|
|
self.logger.warning(f"Critical memory usage: {memory_usage / (1024*1024):.1f} MB")
|
|
self.force_garbage_collection()
|
|
|
|
# Check again after GC
|
|
memory_info = process.memory_info()
|
|
memory_usage = memory_info.rss
|
|
|
|
if memory_usage > self.critical_threshold:
|
|
self.logger.error("Memory usage still critical after garbage collection")
|
|
|
|
elif memory_usage > self.warning_threshold:
|
|
self.logger.info(f"Memory usage warning: {memory_usage / (1024*1024):.1f} MB")
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to check memory usage: {e}")
|
|
|
|
def force_garbage_collection(self):
|
|
"""Force garbage collection to free memory."""
|
|
self.logger.debug("Forcing garbage collection")
|
|
collected = gc.collect()
|
|
self.logger.debug(f"Garbage collection freed {collected} objects")
|
|
|
|
def get_memory_stats(self) -> Dict[str, Any]:
|
|
"""Get current memory statistics."""
|
|
try:
|
|
process = psutil.Process()
|
|
memory_info = process.memory_info()
|
|
|
|
return {
|
|
'rss_mb': memory_info.rss / (1024 * 1024),
|
|
'vms_mb': memory_info.vms / (1024 * 1024),
|
|
'percent': process.memory_percent(),
|
|
'warning_threshold_mb': self.warning_threshold / (1024 * 1024),
|
|
'critical_threshold_mb': self.critical_threshold / (1024 * 1024)
|
|
}
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to get memory stats: {e}")
|
|
return {}
|
|
|
|
|
|
class ParallelDownloadManager:
|
|
"""Manage parallel downloads with configurable thread count."""
|
|
|
|
def __init__(self, max_workers: int = 3, speed_limiter: Optional[SpeedLimiter] = None):
|
|
self.max_workers = max_workers
|
|
self.speed_limiter = speed_limiter or SpeedLimiter()
|
|
self.executor = ThreadPoolExecutor(max_workers=max_workers)
|
|
self.active_tasks: Dict[str, DownloadTask] = {}
|
|
self.pending_queue = queue.PriorityQueue()
|
|
self.completed_tasks: List[DownloadTask] = []
|
|
self.failed_tasks: List[DownloadTask] = []
|
|
self.lock = threading.Lock()
|
|
self.logger = logging.getLogger(__name__)
|
|
self.running = False
|
|
self.worker_thread = None
|
|
|
|
# Statistics
|
|
self.stats = {
|
|
'total_tasks': 0,
|
|
'completed_tasks': 0,
|
|
'failed_tasks': 0,
|
|
'active_tasks': 0,
|
|
'average_speed_mbps': 0.0
|
|
}
|
|
|
|
def start(self):
|
|
"""Start the download manager."""
|
|
if self.running:
|
|
return
|
|
|
|
self.running = True
|
|
self.worker_thread = threading.Thread(target=self._worker_loop, daemon=True)
|
|
self.worker_thread.start()
|
|
self.logger.info(f"Download manager started with {self.max_workers} workers")
|
|
|
|
def stop(self):
|
|
"""Stop the download manager."""
|
|
self.running = False
|
|
|
|
# Cancel all pending tasks
|
|
with self.lock:
|
|
while not self.pending_queue.empty():
|
|
try:
|
|
_, task = self.pending_queue.get_nowait()
|
|
task.status = 'cancelled'
|
|
except queue.Empty:
|
|
break
|
|
|
|
# Shutdown executor
|
|
self.executor.shutdown(wait=True)
|
|
|
|
if self.worker_thread:
|
|
self.worker_thread.join(timeout=5)
|
|
|
|
self.logger.info("Download manager stopped")
|
|
|
|
def add_task(self, task: DownloadTask) -> str:
|
|
"""Add a download task to the queue."""
|
|
with self.lock:
|
|
self.stats['total_tasks'] += 1
|
|
# Priority queue uses negative priority for max-heap behavior
|
|
self.pending_queue.put((-task.priority, task))
|
|
self.logger.info(f"Added download task: {task.task_id}")
|
|
return task.task_id
|
|
|
|
def _worker_loop(self):
|
|
"""Main worker loop that processes download tasks."""
|
|
while self.running:
|
|
try:
|
|
# Check for pending tasks
|
|
if not self.pending_queue.empty() and len(self.active_tasks) < self.max_workers:
|
|
_, task = self.pending_queue.get_nowait()
|
|
|
|
if task.status == 'pending':
|
|
self._start_task(task)
|
|
|
|
# Check completed tasks
|
|
self._check_completed_tasks()
|
|
|
|
time.sleep(0.1) # Small delay to prevent busy waiting
|
|
|
|
except queue.Empty:
|
|
time.sleep(1)
|
|
except Exception as e:
|
|
self.logger.error(f"Error in worker loop: {e}")
|
|
time.sleep(1)
|
|
|
|
def _start_task(self, task: DownloadTask):
|
|
"""Start a download task."""
|
|
with self.lock:
|
|
task.status = 'downloading'
|
|
task.started_at = datetime.now()
|
|
self.active_tasks[task.task_id] = task
|
|
self.stats['active_tasks'] = len(self.active_tasks)
|
|
|
|
# Submit to thread pool
|
|
future = self.executor.submit(self._execute_download, task)
|
|
task.future = future
|
|
|
|
self.logger.info(f"Started download task: {task.task_id}")
|
|
|
|
def _execute_download(self, task: DownloadTask) -> bool:
|
|
"""Execute the actual download."""
|
|
try:
|
|
self.logger.info(f"Executing download: {task.serie_name} S{task.season}E{task.episode}")
|
|
|
|
# Create progress callback that respects speed limiting
|
|
def progress_callback(info):
|
|
if 'downloaded_bytes' in info:
|
|
self.speed_limiter.update_progress(info.get('downloaded_bytes', 0))
|
|
|
|
# Update task progress
|
|
task.progress.update(info)
|
|
|
|
self.speed_limiter.start_download()
|
|
|
|
# Here you would call the actual download function
|
|
# For now, simulate download
|
|
success = self._simulate_download(task, progress_callback)
|
|
|
|
return success
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Download failed for task {task.task_id}: {e}")
|
|
task.error_message = str(e)
|
|
return False
|
|
|
|
def _simulate_download(self, task: DownloadTask, progress_callback: Callable) -> bool:
|
|
"""Simulate download for testing purposes."""
|
|
# This is a placeholder - replace with actual download logic
|
|
total_size = 100 * 1024 * 1024 # 100MB simulation
|
|
downloaded = 0
|
|
chunk_size = 1024 * 1024 # 1MB chunks
|
|
|
|
while downloaded < total_size and task.status == 'downloading':
|
|
# Simulate download chunk
|
|
time.sleep(0.1)
|
|
downloaded += chunk_size
|
|
|
|
progress_info = {
|
|
'status': 'downloading',
|
|
'downloaded_bytes': downloaded,
|
|
'total_bytes': total_size,
|
|
'percent': (downloaded / total_size) * 100
|
|
}
|
|
|
|
progress_callback(progress_info)
|
|
|
|
if downloaded >= total_size:
|
|
progress_callback({'status': 'finished'})
|
|
return True
|
|
|
|
return False
|
|
|
|
def _check_completed_tasks(self):
|
|
"""Check for completed download tasks."""
|
|
completed_task_ids = []
|
|
|
|
with self.lock:
|
|
for task_id, task in self.active_tasks.items():
|
|
if hasattr(task, 'future') and task.future.done():
|
|
completed_task_ids.append(task_id)
|
|
|
|
# Process completed tasks
|
|
for task_id in completed_task_ids:
|
|
self._handle_completed_task(task_id)
|
|
|
|
def _handle_completed_task(self, task_id: str):
|
|
"""Handle a completed download task."""
|
|
with self.lock:
|
|
task = self.active_tasks.pop(task_id, None)
|
|
if not task:
|
|
return
|
|
|
|
task.completed_at = datetime.now()
|
|
self.stats['active_tasks'] = len(self.active_tasks)
|
|
|
|
try:
|
|
success = task.future.result()
|
|
|
|
if success:
|
|
task.status = 'completed'
|
|
self.completed_tasks.append(task)
|
|
self.stats['completed_tasks'] += 1
|
|
self.logger.info(f"Task completed successfully: {task_id}")
|
|
else:
|
|
task.status = 'failed'
|
|
self.failed_tasks.append(task)
|
|
self.stats['failed_tasks'] += 1
|
|
self.logger.warning(f"Task failed: {task_id}")
|
|
|
|
except Exception as e:
|
|
task.status = 'failed'
|
|
task.error_message = str(e)
|
|
self.failed_tasks.append(task)
|
|
self.stats['failed_tasks'] += 1
|
|
self.logger.error(f"Task failed with exception: {task_id} - {e}")
|
|
|
|
def get_task_status(self, task_id: str) -> Optional[Dict[str, Any]]:
|
|
"""Get status of a specific task."""
|
|
with self.lock:
|
|
# Check active tasks
|
|
if task_id in self.active_tasks:
|
|
task = self.active_tasks[task_id]
|
|
return self._task_to_dict(task)
|
|
|
|
# Check completed tasks
|
|
for task in self.completed_tasks:
|
|
if task.task_id == task_id:
|
|
return self._task_to_dict(task)
|
|
|
|
# Check failed tasks
|
|
for task in self.failed_tasks:
|
|
if task.task_id == task_id:
|
|
return self._task_to_dict(task)
|
|
|
|
return None
|
|
|
|
def _task_to_dict(self, task: DownloadTask) -> Dict[str, Any]:
|
|
"""Convert task to dictionary representation."""
|
|
return {
|
|
'task_id': task.task_id,
|
|
'serie_name': task.serie_name,
|
|
'season': task.season,
|
|
'episode': task.episode,
|
|
'status': task.status,
|
|
'progress': task.progress,
|
|
'created_at': task.created_at.isoformat(),
|
|
'started_at': task.started_at.isoformat() if task.started_at else None,
|
|
'completed_at': task.completed_at.isoformat() if task.completed_at else None,
|
|
'error_message': task.error_message,
|
|
'retry_count': task.retry_count
|
|
}
|
|
|
|
def get_all_tasks(self) -> Dict[str, List[Dict[str, Any]]]:
|
|
"""Get all tasks grouped by status."""
|
|
with self.lock:
|
|
return {
|
|
'active': [self._task_to_dict(task) for task in self.active_tasks.values()],
|
|
'completed': [self._task_to_dict(task) for task in self.completed_tasks[-50:]], # Last 50
|
|
'failed': [self._task_to_dict(task) for task in self.failed_tasks[-50:]] # Last 50
|
|
}
|
|
|
|
def get_statistics(self) -> Dict[str, Any]:
|
|
"""Get download manager statistics."""
|
|
return self.stats.copy()
|
|
|
|
def set_max_workers(self, max_workers: int):
|
|
"""Change the number of worker threads."""
|
|
if max_workers <= 0:
|
|
raise ValueError("max_workers must be positive")
|
|
|
|
self.max_workers = max_workers
|
|
|
|
# Recreate executor with new worker count
|
|
old_executor = self.executor
|
|
self.executor = ThreadPoolExecutor(max_workers=max_workers)
|
|
old_executor.shutdown(wait=False)
|
|
|
|
self.logger.info(f"Updated worker count to {max_workers}")
|
|
|
|
|
|
class ResumeManager:
|
|
"""Manage download resumption for interrupted downloads."""
|
|
|
|
def __init__(self, resume_dir: str = "./resume"):
|
|
self.resume_dir = resume_dir
|
|
self.logger = logging.getLogger(__name__)
|
|
os.makedirs(resume_dir, exist_ok=True)
|
|
|
|
def save_resume_info(self, task_id: str, resume_data: Dict[str, Any]):
|
|
"""Save resume information for a download."""
|
|
try:
|
|
resume_file = os.path.join(self.resume_dir, f"{task_id}.json")
|
|
with open(resume_file, 'w') as f:
|
|
json.dump(resume_data, f, indent=2, default=str)
|
|
|
|
self.logger.debug(f"Saved resume info for task: {task_id}")
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to save resume info for {task_id}: {e}")
|
|
|
|
def load_resume_info(self, task_id: str) -> Optional[Dict[str, Any]]:
|
|
"""Load resume information for a download."""
|
|
try:
|
|
resume_file = os.path.join(self.resume_dir, f"{task_id}.json")
|
|
|
|
if os.path.exists(resume_file):
|
|
with open(resume_file, 'r') as f:
|
|
resume_data = json.load(f)
|
|
|
|
self.logger.debug(f"Loaded resume info for task: {task_id}")
|
|
return resume_data
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to load resume info for {task_id}: {e}")
|
|
|
|
return None
|
|
|
|
def clear_resume_info(self, task_id: str):
|
|
"""Clear resume information after successful completion."""
|
|
try:
|
|
resume_file = os.path.join(self.resume_dir, f"{task_id}.json")
|
|
|
|
if os.path.exists(resume_file):
|
|
os.remove(resume_file)
|
|
self.logger.debug(f"Cleared resume info for task: {task_id}")
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to clear resume info for {task_id}: {e}")
|
|
|
|
def get_resumable_tasks(self) -> List[str]:
|
|
"""Get list of tasks that can be resumed."""
|
|
try:
|
|
resume_files = [f for f in os.listdir(self.resume_dir) if f.endswith('.json')]
|
|
task_ids = [os.path.splitext(f)[0] for f in resume_files]
|
|
return task_ids
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to get resumable tasks: {e}")
|
|
return []
|
|
|
|
|
|
# Global instances
|
|
speed_limiter = SpeedLimiter()
|
|
memory_monitor = MemoryMonitor()
|
|
download_manager = ParallelDownloadManager(max_workers=3, speed_limiter=speed_limiter)
|
|
resume_manager = ResumeManager()
|
|
|
|
|
|
def init_performance_monitoring():
|
|
"""Initialize performance monitoring components."""
|
|
memory_monitor.start_monitoring()
|
|
download_manager.start()
|
|
|
|
|
|
def cleanup_performance_monitoring():
|
|
"""Clean up performance monitoring components."""
|
|
memory_monitor.stop_monitoring()
|
|
download_manager.stop()
|
|
|
|
|
|
# Export main components
|
|
__all__ = [
|
|
'SpeedLimiter',
|
|
'MemoryMonitor',
|
|
'ParallelDownloadManager',
|
|
'ResumeManager',
|
|
'DownloadTask',
|
|
'speed_limiter',
|
|
'download_cache',
|
|
'memory_monitor',
|
|
'download_manager',
|
|
'resume_manager',
|
|
'init_performance_monitoring',
|
|
'cleanup_performance_monitoring'
|
|
] |