Add hybrid caching system for better performance & fix limitation on multi workers

Replaced TTLCache with an optimized hybrid cache combining memory and file storage for better performance and scalability. Removed cachetools dependency and updated pyproject.toml to include aiofiles. Adjusted speed test service to use the new caching mechanism.

mediaflow_proxy/speedtest/service.py

@@ -1,17 +1,17 @@
 import logging
 import time
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Dict, Optional, Type
 
-from cachetools import TTLCache
 from httpx import AsyncClient
 
+from mediaflow_proxy.configs import settings
 from mediaflow_proxy.utils.http_utils import Streamer
+from mediaflow_proxy.utils.cache_utils import get_cached_speedtest, set_cache_speedtest
 from .models import SpeedTestTask, LocationResult, SpeedTestResult, SpeedTestProvider
 from .providers.all_debrid import AllDebridSpeedTest
 from .providers.base import BaseSpeedTestProvider
 from .providers.real_debrid import RealDebridSpeedTest
-from ..configs import settings
 
 logger = logging.getLogger(__name__)
 
@@ -20,9 +20,6 @@ class SpeedTestService:
     """Service for managing speed tests across different providers."""
 
     def __init__(self):
-        # Cache for speed test results (1 hour TTL)
-        self._cache: TTLCache[str, SpeedTestTask] = TTLCache(maxsize=100, ttl=3600)
-
         # Provider mapping
         self._providers: Dict[SpeedTestProvider, Type[BaseSpeedTestProvider]] = {
             SpeedTestProvider.REAL_DEBRID: RealDebridSpeedTest,
@@ -49,19 +46,22 @@ class SpeedTestService:
         # Get initial URLs and user info
         urls, user_info = await provider_impl.get_test_urls()
 
-        task = SpeedTestTask(task_id=task_id, provider=provider, started_at=datetime.utcnow(), user_info=user_info)
+        task = SpeedTestTask(
+            task_id=task_id, provider=provider, started_at=datetime.now(tz=timezone.utc), user_info=user_info
+        )
 
-        self._cache[task_id] = task
+        await set_cache_speedtest(task_id, task)
         return task
 
-    async def get_test_results(self, task_id: str) -> Optional[SpeedTestTask]:
+    @staticmethod
+    async def get_test_results(task_id: str) -> Optional[SpeedTestTask]:
         """Get results for a specific task."""
-        return self._cache.get(task_id)
+        return await get_cached_speedtest(task_id)
 
     async def run_speedtest(self, task_id: str, provider: SpeedTestProvider, api_key: Optional[str] = None):
         """Run the speed test with real-time updates."""
         try:
-            task = self._cache.get(task_id)
+            task = await get_cached_speedtest(task_id)
             if not task:
                 raise ValueError(f"Task {task_id} not found")
 
@@ -74,27 +74,28 @@ class SpeedTestService:
                 for location, url in config.test_urls.items():
                     try:
                         task.current_location = location
+                        await set_cache_speedtest(task_id, task)
                         result = await self._test_location(location, url, streamer, config.test_duration, provider_impl)
                         task.results[location] = result
-                        self._cache[task_id] = task
+                        await set_cache_speedtest(task_id, task)
                     except Exception as e:
                         logger.error(f"Error testing {location}: {str(e)}")
                         task.results[location] = LocationResult(
                             error=str(e), server_name=location, server_url=config.test_urls[location]
                         )
-                        self._cache[task_id] = task
+                        await set_cache_speedtest(task_id, task)
 
             # Mark task as completed
-            task.completed_at = datetime.utcnow()
+            task.completed_at = datetime.now(tz=timezone.utc)
             task.status = "completed"
             task.current_location = None
-            self._cache[task_id] = task
+            await set_cache_speedtest(task_id, task)
 
         except Exception as e:
             logger.error(f"Error in speed test task {task_id}: {str(e)}")
-            if task := self._cache.get(task_id):
+            if task := await get_cached_speedtest(task_id):
                 task.status = "failed"
-                self._cache[task_id] = task
+                await set_cache_speedtest(task_id, task)
 
     async def _test_location(
         self, location: str, url: str, streamer: Streamer, test_duration: int, provider: BaseSpeedTestProvider

mediaflow_proxy/utils/cache_utils.py

@@ -1,16 +1,346 @@
-import datetime
+import asyncio
+import json
 import logging
+import os
+import tempfile
+import threading
+import time
+from collections import OrderedDict
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Union, Any
 
-from cachetools import TTLCache
+import aiofiles
+import aiofiles.os
+from pydantic import ValidationError
 
-from .http_utils import download_file_with_retry
-from .mpd_utils import parse_mpd, parse_mpd_dict
+from mediaflow_proxy.speedtest.models import SpeedTestTask
+from mediaflow_proxy.utils.http_utils import download_file_with_retry
+from mediaflow_proxy.utils.mpd_utils import parse_mpd, parse_mpd_dict
 
 logger = logging.getLogger(__name__)
 
-# cache dictionary
-mpd_cache = TTLCache(maxsize=100, ttl=300)  # 5 minutes default TTL
-init_segment_cache = TTLCache(maxsize=100, ttl=3600)  # 1 hour default TTL
+
+@dataclass
+class CacheEntry:
+    """Represents a cache entry with metadata."""
+
+    data: bytes
+    expires_at: float
+    access_count: int = 0
+    last_access: float = 0.0
+    size: int = 0
+
+
+class CacheStats:
+    """Tracks cache performance metrics."""
+
+    def __init__(self):
+        self.hits = 0
+        self.misses = 0
+        self.memory_hits = 0
+        self.disk_hits = 0
+        self._lock = threading.Lock()
+
+    def record_hit(self, from_memory: bool):
+        with self._lock:
+            self.hits += 1
+            if from_memory:
+                self.memory_hits += 1
+            else:
+                self.disk_hits += 1
+
+    def record_miss(self):
+        with self._lock:
+            self.misses += 1
+
+    @property
+    def hit_rate(self) -> float:
+        total = self.hits + self.misses
+        return self.hits / total if total > 0 else 0.0
+
+    def __str__(self) -> str:
+        return (
+            f"Cache Stats: Hits={self.hits} (Memory: {self.memory_hits}, "
+            f"Disk: {self.disk_hits}), Misses={self.misses}, "
+            f"Hit Rate={self.hit_rate:.2%}"
+        )
+
+
+class AsyncLRUMemoryCache:
+    """Thread-safe LRU memory cache with async support."""
+
+    def __init__(self, maxsize: int):
+        self.maxsize = maxsize
+        self._cache: OrderedDict[str, CacheEntry] = OrderedDict()
+        self._lock = threading.Lock()
+        self._current_size = 0
+
+    def get(self, key: str) -> Optional[CacheEntry]:
+        with self._lock:
+            if key in self._cache:
+                entry = self._cache.pop(key)  # Remove and re-insert for LRU
+                if time.time() < entry.expires_at:
+                    entry.access_count += 1
+                    entry.last_access = time.time()
+                    self._cache[key] = entry
+                    return entry
+                else:
+                    # Remove expired entry
+                    self._current_size -= entry.size
+                    del self._cache[key]
+            return None
+
+    def set(self, key: str, entry: CacheEntry) -> None:
+        with self._lock:
+            if key in self._cache:
+                old_entry = self._cache[key]
+                self._current_size -= old_entry.size
+
+            # Check if we need to make space
+            while self._current_size + entry.size > self.maxsize and self._cache:
+                _, removed_entry = self._cache.popitem(last=False)
+                self._current_size -= removed_entry.size
+
+            self._cache[key] = entry
+            self._current_size += entry.size
+
+    def remove(self, key: str) -> None:
+        with self._lock:
+            if key in self._cache:
+                entry = self._cache.pop(key)
+                self._current_size -= entry.size
+
+
+class OptimizedHybridCache:
+    """High-performance hybrid cache combining memory and file storage."""
+
+    def __init__(
+        self,
+        cache_dir_name: str,
+        ttl: int,
+        max_memory_size: int = 100 * 1024 * 1024,  # 100MB default
+        file_shards: int = 256,  # Number of subdirectories for sharding
+        executor_workers: int = 4,
+    ):
+        self.cache_dir = Path(tempfile.gettempdir()) / cache_dir_name
+        self.ttl = ttl
+        self.file_shards = file_shards
+        self.memory_cache = AsyncLRUMemoryCache(maxsize=max_memory_size)
+        self.stats = CacheStats()
+        self._executor = ThreadPoolExecutor(max_workers=executor_workers)
+        self._lock = asyncio.Lock()
+
+        # Initialize cache directories
+        self._init_cache_dirs()
+
+    def _init_cache_dirs(self):
+        """Initialize sharded cache directories."""
+        for i in range(self.file_shards):
+            shard_dir = self.cache_dir / f"shard_{i:03d}"
+            os.makedirs(shard_dir, exist_ok=True)
+
+    def _get_shard_path(self, key: str) -> Path:
+        """Get the appropriate shard directory for a key."""
+        shard_num = hash(key) % self.file_shards
+        return self.cache_dir / f"shard_{shard_num:03d}"
+
+    def _get_file_path(self, key: str) -> Path:
+        """Get the file path for a cache key."""
+        safe_key = str(hash(key))
+        return self._get_shard_path(key) / safe_key
+
+    async def get(self, key: str, default: Any = None) -> Optional[bytes]:
+        """
+        Get value from cache, trying memory first then file.
+
+        Args:
+            key: Cache key
+            default: Default value if key not found
+
+        Returns:
+            Cached value or default if not found
+        """
+        # Try memory cache first
+        entry = self.memory_cache.get(key)
+        if entry is not None:
+            self.stats.record_hit(from_memory=True)
+            return entry.data
+
+        # Try file cache
+        try:
+            file_path = self._get_file_path(key)
+            async with aiofiles.open(file_path, "rb") as f:
+                metadata_size = await f.read(8)
+                metadata_length = int.from_bytes(metadata_size, "big")
+                metadata_bytes = await f.read(metadata_length)
+                metadata = json.loads(metadata_bytes.decode())
+
+                # Check expiration
+                if metadata["expires_at"] < time.time():
+                    await self.delete(key)
+                    self.stats.record_miss()
+                    return default
+
+                # Read data
+                data = await f.read()
+
+                # Update memory cache in background
+                entry = CacheEntry(
+                    data=data,
+                    expires_at=metadata["expires_at"],
+                    access_count=metadata["access_count"] + 1,
+                    last_access=time.time(),
+                    size=len(data),
+                )
+                self.memory_cache.set(key, entry)
+
+                self.stats.record_hit(from_memory=False)
+                return data
+
+        except FileNotFoundError:
+            self.stats.record_miss()
+            return default
+        except Exception as e:
+            logger.error(f"Error reading from cache: {e}")
+            self.stats.record_miss()
+            return default
+
+    async def set(self, key: str, data: Union[bytes, bytearray, memoryview], ttl: Optional[int] = None) -> bool:
+        """
+        Set value in both memory and file cache.
+
+        Args:
+            key: Cache key
+            data: Data to cache
+            ttl: Optional TTL override
+
+        Returns:
+            bool: Success status
+        """
+        if not isinstance(data, (bytes, bytearray, memoryview)):
+            raise ValueError("Data must be bytes, bytearray, or memoryview")
+
+        expires_at = time.time() + (ttl or self.ttl)
+
+        # Create cache entry
+        entry = CacheEntry(data=data, expires_at=expires_at, access_count=0, last_access=time.time(), size=len(data))
+
+        # Update memory cache
+        self.memory_cache.set(key, entry)
+
+        # Update file cache
+        try:
+            file_path = self._get_file_path(key)
+            temp_path = file_path.with_suffix(".tmp")
+
+            metadata = {"expires_at": expires_at, "access_count": 0, "last_access": time.time()}
+            metadata_bytes = json.dumps(metadata).encode()
+            metadata_size = len(metadata_bytes).to_bytes(8, "big")
+
+            async with aiofiles.open(temp_path, "wb") as f:
+                await f.write(metadata_size)
+                await f.write(metadata_bytes)
+                await f.write(data)
+
+            await aiofiles.os.rename(temp_path, file_path)
+            return True
+
+        except Exception as e:
+            logger.error(f"Error writing to cache: {e}")
+            try:
+                await aiofiles.os.remove(temp_path)
+            except:
+                pass
+            return False
+
+    async def delete(self, key: str) -> bool:
+        """Delete item from both caches."""
+        self.memory_cache.remove(key)
+
+        try:
+            file_path = self._get_file_path(key)
+            await aiofiles.os.remove(file_path)
+            return True
+        except FileNotFoundError:
+            return True
+        except Exception as e:
+            logger.error(f"Error deleting from cache: {e}")
+            return False
+
+    async def cleanup_expired(self):
+        """Clean up expired cache entries."""
+        current_time = time.time()
+
+        async def check_and_clean_file(file_path: Path):
+            try:
+                async with aiofiles.open(file_path, "rb") as f:
+                    metadata_size = await f.read(8)
+                    metadata_length = int.from_bytes(metadata_size, "big")
+                    metadata_bytes = await f.read(metadata_length)
+                    metadata = json.loads(metadata_bytes.decode())
+
+                    if metadata["expires_at"] < current_time:
+                        await aiofiles.os.remove(file_path)
+            except Exception as e:
+                logger.error(f"Error cleaning up file {file_path}: {e}")
+
+        # Clean up each shard
+        for i in range(self.file_shards):
+            shard_dir = self.cache_dir / f"shard_{i:03d}"
+            try:
+                async for entry in aiofiles.os.scandir(shard_dir):
+                    if entry.is_file() and not entry.name.endswith(".tmp"):
+                        await check_and_clean_file(Path(entry.path))
+            except Exception as e:
+                logger.error(f"Error scanning shard directory {shard_dir}: {e}")
+
+
+# Create cache instances
+INIT_SEGMENT_CACHE = OptimizedHybridCache(
+    cache_dir_name="init_segment_cache",
+    ttl=3600,  # 1 hour
+    max_memory_size=500 * 1024 * 1024,  # 500MB for init segments
+    file_shards=512,  # More shards for better distribution
+)
+
+MPD_CACHE = OptimizedHybridCache(
+    cache_dir_name="mpd_cache",
+    ttl=300,  # 5 minutes
+    max_memory_size=100 * 1024 * 1024,  # 100MB for MPD files
+    file_shards=128,
+)
+
+SPEEDTEST_CACHE = OptimizedHybridCache(
+    cache_dir_name="speedtest_cache",
+    ttl=3600,  # 1 hour
+    max_memory_size=50 * 1024 * 1024,  # 50MB for speed test results
+    file_shards=64,
+)
+
+
+# Specific cache implementations
+async def get_cached_init_segment(
+    init_url: str, headers: dict, verify_ssl: bool = True, use_request_proxy: bool = True
+) -> Optional[bytes]:
+    """Get initialization segment from cache or download it."""
+    # Try cache first
+    cached_data = await INIT_SEGMENT_CACHE.get(init_url)
+    if cached_data is not None:
+        return cached_data
+
+    # Download if not cached
+    try:
+        init_content = await download_file_with_retry(
+            init_url, headers, verify_ssl=verify_ssl, use_request_proxy=use_request_proxy
+        )
+        if init_content:
+            await INIT_SEGMENT_CACHE.set(init_url, init_content)
+        return init_content
+    except Exception as e:
+        logger.error(f"Error downloading init segment: {e}")
+        return None
 
 
 async def get_cached_mpd(
@@ -20,54 +350,49 @@ async def get_cached_mpd(
     parse_segment_profile_id: str | None = None,
     verify_ssl: bool = True,
     use_request_proxy: bool = True,
-) -> dict:
-    """
-    Retrieves and caches the MPD manifest, parsing it if not already cached.
-
-    Args:
-        mpd_url (str): The URL of the MPD manifest.
-        headers (dict): The headers to include in the request.
-        parse_drm (bool): Whether to parse DRM information.
-        parse_segment_profile_id (str, optional): The profile ID to parse segments for. Defaults to None.
-        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.
-        use_request_proxy (bool, optional): Whether to use the proxy configuration from the user's MediaFlow config. Defaults to True.
-
-    Returns:
-        dict: The parsed MPD manifest data.
-    """
-    current_time = datetime.datetime.now(datetime.UTC)
-    if mpd_url in mpd_cache and mpd_cache[mpd_url]["expires"] > current_time:
-        logger.info(f"Using cached MPD for {mpd_url}")
-        return parse_mpd_dict(mpd_cache[mpd_url]["mpd"], mpd_url, parse_drm, parse_segment_profile_id)
-
-    mpd_dict = parse_mpd(
-        await download_file_with_retry(mpd_url, headers, verify_ssl=verify_ssl, use_request_proxy=use_request_proxy)
-    )
-    parsed_mpd_dict = parse_mpd_dict(mpd_dict, mpd_url, parse_drm, parse_segment_profile_id)
-    current_time = datetime.datetime.now(datetime.UTC)
-    expiration_time = current_time + datetime.timedelta(seconds=parsed_mpd_dict.get("minimumUpdatePeriod", 300))
-    mpd_cache[mpd_url] = {"mpd": mpd_dict, "expires": expiration_time}
-    return parsed_mpd_dict
-
+) -> Optional[dict]:
+    """Get MPD from cache or download and parse it."""
+    # Try cache first
+    cached_data = await MPD_CACHE.get(mpd_url)
+    if cached_data is not None:
+        try:
+            mpd_dict = json.loads(cached_data)
+            return parse_mpd_dict(mpd_dict, mpd_url, parse_drm, parse_segment_profile_id)
+        except json.JSONDecodeError:
+            await MPD_CACHE.delete(mpd_url)
 
-async def get_cached_init_segment(
-    init_url: str, headers: dict, verify_ssl: bool = True, use_request_proxy: bool = True
-) -> bytes:
-    """
-    Retrieves and caches the initialization segment.
-
-    Args:
-        init_url (str): The URL of the initialization segment.
-        headers (dict): The headers to include in the request.
-        verify_ssl (bool, optional): Whether to verify the SSL certificate of the destination. Defaults to True.
-        use_request_proxy (bool, optional): Whether to use the proxy configuration from the user's MediaFlow config. Defaults to True.
-
-    Returns:
-        bytes: The initialization segment content.
-    """
-    if init_url not in init_segment_cache:
-        init_content = await download_file_with_retry(
-            init_url, headers, verify_ssl=verify_ssl, use_request_proxy=use_request_proxy
+    # Download and parse if not cached
+    try:
+        mpd_content = await download_file_with_retry(
+            mpd_url, headers, verify_ssl=verify_ssl, use_request_proxy=use_request_proxy
         )
-        init_segment_cache[init_url] = init_content
-    return init_segment_cache[init_url]
+        mpd_dict = parse_mpd(mpd_content)
+        parsed_dict = parse_mpd_dict(mpd_dict, mpd_url, parse_drm, parse_segment_profile_id)
+
+        # Cache the original MPD dict
+        await MPD_CACHE.set(mpd_url, json.dumps(mpd_dict).encode())
+        return parsed_dict
+    except Exception as e:
+        logger.error(f"Error processing MPD: {e}")
+        return None
+
+
+async def get_cached_speedtest(task_id: str) -> Optional[SpeedTestTask]:
+    """Get speed test results from cache."""
+    cached_data = await SPEEDTEST_CACHE.get(task_id)
+    if cached_data is not None:
+        try:
+            return SpeedTestTask.model_validate_json(cached_data.decode())
+        except ValidationError as e:
+            logger.error(f"Error parsing cached speed test data: {e}")
+            await SPEEDTEST_CACHE.delete(task_id)
+    return None
+
+
+async def set_cache_speedtest(task_id: str, task: SpeedTestTask) -> bool:
+    """Cache speed test results."""
+    try:
+        return await SPEEDTEST_CACHE.set(task_id, task.model_dump_json().encode())
+    except Exception as e:
+        logger.error(f"Error caching speed test data: {e}")
+        return False

poetry.lock

@@ -1,5 +1,16 @@
 # This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand.
 
+[[package]]
+name = "aiofiles"
+version = "24.1.0"
+description = "File support for asyncio."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "aiofiles-24.1.0-py3-none-any.whl", hash = "sha256:b4ec55f4195e3eb5d7abd1bf7e061763e864dd4954231fb8539a0ef8bb8260e5"},
+    {file = "aiofiles-24.1.0.tar.gz", hash = "sha256:22a075c9e5a3810f0c2e48f3008c94d68c65d763b9b03857924c99e57355166c"},
+]
+
 [[package]]
 name = "annotated-types"
 version = "0.7.0"
@@ -79,17 +90,6 @@ d = ["aiohttp (>=3.10)"]
 jupyter = ["ipython (>=7.8.0)", "tokenize-rt (>=3.2.0)"]
 uvloop = ["uvloop (>=0.15.2)"]
 
-[[package]]
-name = "cachetools"
-version = "5.5.0"
-description = "Extensible memoizing collections and decorators"
-optional = false
-python-versions = ">=3.7"
-files = [
-    {file = "cachetools-5.5.0-py3-none-any.whl", hash = "sha256:02134e8439cdc2ffb62023ce1debca2944c3f289d66bb17ead3ab3dede74b292"},
-    {file = "cachetools-5.5.0.tar.gz", hash = "sha256:2cc24fb4cbe39633fb7badd9db9ca6295d766d9c2995f245725a46715d050f2a"},
-]
-
 [[package]]
 name = "certifi"
 version = "2024.8.30"
@@ -631,4 +631,4 @@ files = [
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.10"
-content-hash = "b9b0a0539cd08dd58b46d7fc41e940666ae658a68384c19a2fb692001aa71120"
+content-hash = "31ae30f007ef7dacc5a13f41d1f88f3d2112e10e72846df59cb0956593bb33b9"

pyproject.toml

@@ -27,12 +27,12 @@ fastapi = "0.115.4"
 httpx = {extras = ["socks"], version = "^0.27.2"}
 tenacity = "^9.0.0"
 xmltodict = "^0.14.2"
-cachetools = "^5.4.0"
 pydantic-settings = "^2.6.1"
 gunicorn = "^23.0.0"
 pycryptodome = "^3.20.0"
 uvicorn = "^0.32.0"
 tqdm = "^4.67.0"
+aiofiles = "^24.1.0"
 
 
 [tool.poetry.group.dev.dependencies]