Merge pull request #415 from g0ldyy/feat/better-performances

feat: improve performance with process pool executor and optimize torrent caching/database writes

comet/api/app.py

@@ -16,6 +16,7 @@ from comet.background_scraper.worker import background_scraper
 from comet.core.database import (cleanup_expired_locks,
                                  cleanup_expired_sessions, setup_database,
                                  teardown_database)
+from comet.core.execution import setup_executor, shutdown_executor
 from comet.core.logger import logger
 from comet.core.models import settings
 from comet.services.anime import anime_mapper
@@ -49,6 +50,7 @@ async def lifespan(app: FastAPI):
     # loop.set_debug(True)
 
     await setup_database()
+    setup_executor()
     await download_best_trackers()
 
     # Load anime ID mapping for enhanced metadata and anime detection
@@ -107,6 +109,7 @@ async def lifespan(app: FastAPI):
         await torrent_update_queue.stop()
 
         await teardown_database()
+        shutdown_executor()
 
 
 tags_metadata = [

comet/api/endpoints/stream.py

@@ -8,6 +8,7 @@ from fastapi import APIRouter, BackgroundTasks, Request
 from comet.core.config_validation import config_check
 from comet.core.logger import logger
 from comet.core.models import database, settings, trackers
+from comet.debrid.exceptions import DebridAuthError
 from comet.debrid.manager import get_debrid_extension
 from comet.metadata.filter import release_filter
 from comet.metadata.manager import MetadataScraper
@@ -413,14 +414,25 @@ async def stream(
             and debrid_service != "torrent"
         ):
             logger.log("SCRAPER", "🔄 Checking availability on debrid service...")
-            await debrid_service_instance.get_and_cache_availability(
-                session,
-                torrent_manager.torrents,
-                media_id,
-                media_only_id,
-                season,
-                episode,
-            )
+            try:
+                await debrid_service_instance.get_and_cache_availability(
+                    session,
+                    torrent_manager.torrents,
+                    media_id,
+                    media_only_id,
+                    season,
+                    episode,
+                )
+            except DebridAuthError as e:
+                return {
+                    "streams": [
+                        {
+                            "name": "[❌] Comet",
+                            "description": e.display_message,
+                            "url": "https://comet.fast",
+                        }
+                    ]
+                }
 
         if debrid_service != "torrent":
             cached_count = sum(

comet/core/execution.py

@@ -0,0 +1,18 @@
+from concurrent.futures import ProcessPoolExecutor
+
+app_executor = None
+
+
+def setup_executor():
+    global app_executor
+    app_executor = ProcessPoolExecutor()
+
+
+def shutdown_executor():
+    global app_executor
+    if app_executor:
+        app_executor.shutdown()
+
+
+def get_executor():
+    return app_executor

comet/debrid/exceptions.py

@@ -0,0 +1,20 @@
+class DebridError(Exception):
+    """Base exception for debrid-related errors."""
+
+    def __init__(self, message: str, display_message: str = None):
+        self.message = message
+        self.display_message = display_message or message
+        super().__init__(self.message)
+
+
+class DebridAuthError(DebridError):
+    """Raised when debrid authentication fails (not premium, invalid API key, etc.)."""
+
+    def __init__(self, debrid_name: str, message: str = None):
+        self.debrid_name = debrid_name
+        default_message = f"{debrid_name}: Authentication failed or not premium"
+        display_message = (
+            message
+            or f"{debrid_name}: Invalid API key or no active subscription.\nPlease check your debrid account."
+        )
+        super().__init__(default_message, display_message)

comet/debrid/stremthru.py

@@ -4,13 +4,19 @@ from urllib.parse import quote, unquote
 import aiohttp
 from RTN import parse, title_match
 
+from comet.core.execution import get_executor
 from comet.core.logger import logger
 from comet.core.models import settings
+from comet.debrid.exceptions import DebridAuthError
 from comet.services.debrid_cache import cache_availability
 from comet.services.torrent_manager import torrent_update_queue
 from comet.utils.parsing import is_video
 
 
+def batch_parse(filenames):
+    return [parse(f) for f in filenames]
+
+
 class StremThru:
     def __init__(
         self,
@@ -43,17 +49,29 @@ class StremThru:
 
     async def check_premium(self):
         try:
-            user = await self.session.get(
+            response = await self.session.get(
                 f"{self.base_url}/user?client_ip={self.client_ip}"
             )
-            user = await user.json()
-            return user["data"]["subscription_status"] == "premium"
-        except Exception as e:
-            logger.warning(
-                f"Exception while checking premium status on {self.name}: {e}"
-            )
+            user = await response.json()
 
-        return False
+            if "data" not in user:
+                raise DebridAuthError(
+                    self.name,
+                    f"{self.name}: Invalid API key.\nPlease check your configuration.",
+                )
+
+            if user["data"]["subscription_status"] != "premium":
+                raise DebridAuthError(
+                    self.name,
+                    f"{self.name}: No active subscription.\nPlease renew your debrid account.",
+                )
+        except DebridAuthError:
+            raise
+        except Exception as e:
+            raise DebridAuthError(
+                self.name,
+                f"{self.name}: Failed to check account status.\n{e}",
+            )
 
     async def get_instant(self, magnets: list):
         try:
@@ -72,8 +90,7 @@ class StremThru:
         tracker_map: dict,
         sources_map: dict,
     ):
-        if not await self.check_premium():
-            return []
+        await self.check_premium()
 
         chunk_size = 50
         chunks = [
@@ -95,6 +112,26 @@ class StremThru:
 
         is_offcloud = self.real_debrid_name == "offcloud"
 
+        filenames_to_parse = []
+        if not is_offcloud:
+            for result in availability:
+                for torrent in result:
+                    if torrent["status"] != "cached":
+                        continue
+                    for file in torrent["files"]:
+                        filename = file["name"].split("/")[-1]
+                        if not is_video(filename) or "sample" in filename.lower():
+                            continue
+                        filenames_to_parse.append(filename)
+
+        parsed_iter = iter([])
+        if filenames_to_parse:
+            loop = asyncio.get_running_loop()
+            parsed_results = await loop.run_in_executor(
+                get_executor(), batch_parse, filenames_to_parse
+            )
+            parsed_iter = iter(parsed_results)
+
         files = []
         cached_count = 0
         for result in availability:
@@ -127,7 +164,7 @@ class StremThru:
                         if not is_video(filename) or "sample" in filename.lower():
                             continue
 
-                        filename_parsed = parse(filename)
+                        filename_parsed = next(parsed_iter)
 
                         season = (
                             filename_parsed.seasons[0]

comet/main.py

@@ -44,20 +44,18 @@ def run_with_uvicorn():
         workers=settings.FASTAPI_WORKERS,
         log_config=None,
     )
-    server = Server(config=config)
+    server = uvicorn.Server(config=config)
 
-    with server.run_in_thread():
-        log_startup_info(settings)
-        try:
-            while True:
-                time.sleep(1)  # Keep the main thread alive
-        except KeyboardInterrupt:
-            logger.log("COMET", "Server stopped by user")
-        except Exception as e:
-            logger.error(f"Unexpected error: {e}")
-            logger.exception(traceback.format_exc())
-        finally:
-            logger.log("COMET", "Server Shutdown")
+    log_startup_info(settings)
+    try:
+        server.run()
+    except KeyboardInterrupt:
+        logger.log("COMET", "Server stopped by user")
+    except Exception as e:
+        logger.error(f"Unexpected error: {e}")
+        logger.exception(traceback.format_exc())
+    finally:
+        logger.log("COMET", "Server Shutdown")
 
 
 def run_with_gunicorn():

comet/services/orchestration.py

@@ -5,12 +5,13 @@ import aiohttp
 import orjson
 from RTN import DefaultRanking, ParsedData
 
+from comet.core.execution import get_executor
 from comet.core.logger import logger
 from comet.core.models import CometSettingsModel, database, settings
 from comet.scrapers.manager import scraper_manager
 from comet.services.filtering import filter_worker
 from comet.services.ranking import rank_worker
-from comet.utils.parsing import default_dump
+from comet.services.torrent_manager import torrent_update_queue
 
 
 class TorrentManager:
@@ -72,7 +73,7 @@ class TorrentManager:
         async for scraper_name, results in scraper_manager.scrape_all(request, session):
             await self.filter_manager(scraper_name, results)
 
-        await self.cache_torrents()
+        asyncio.create_task(self.cache_torrents())
 
         for torrent in self.ready_to_cache:
             season = torrent["parsed"].seasons[0] if torrent["parsed"].seasons else None
@@ -128,41 +129,24 @@ class TorrentManager:
             }
 
     async def cache_torrents(self):
-        current_time = time.time()
-        values = [
-            {
-                "media_id": self.media_only_id,
+        for torrent in self.ready_to_cache:
+            file_info = {
                 "info_hash": torrent["infoHash"],
-                "file_index": int(torrent["fileIndex"])
-                if torrent["fileIndex"] is not None
-                else None,
+                "index": torrent["fileIndex"],
+                "title": torrent["title"],
+                "size": torrent["size"],
                 "season": torrent["parsed"].seasons[0]
                 if torrent["parsed"].seasons
                 else self.season,
                 "episode": torrent["parsed"].episodes[0]
                 if torrent["parsed"].episodes
                 else None,
-                "title": torrent["title"],
-                "seeders": int(torrent["seeders"])
-                if torrent["seeders"] is not None
-                else None,
-                "size": int(torrent["size"]) if torrent["size"] is not None else None,
+                "parsed": torrent["parsed"],
+                "seeders": torrent["seeders"],
                 "tracker": torrent["tracker"],
-                "sources": orjson.dumps(torrent["sources"]).decode("utf-8"),
-                "parsed": orjson.dumps(torrent["parsed"], default_dump).decode("utf-8"),
-                "timestamp": current_time,
+                "sources": torrent["sources"],
             }
-            for torrent in self.ready_to_cache
-        ]
-
-        query = f"""
-            INSERT {"OR REPLACE " if settings.DATABASE_TYPE == "sqlite" else ""}
-            INTO torrents
-            VALUES (:media_id, :info_hash, :file_index, :season, :episode, :title, :seeders, :size, :tracker, :sources, :parsed, :timestamp)
-            {" ON CONFLICT DO NOTHING" if settings.DATABASE_TYPE == "postgresql" else ""}
-        """
-
-        await database.execute_many(query, values)
+            await torrent_update_queue.add_torrent_info(file_info, self.media_only_id)
 
     async def filter_manager(self, scraper_name: str, torrents: list):
         if len(torrents) == 0:
@@ -191,7 +175,7 @@ class TorrentManager:
         chunk_size = 50
         tasks = [
             loop.run_in_executor(
-                None,
+                get_executor(),
                 filter_worker,
                 new_torrents[i : i + chunk_size],
                 self.title,
@@ -217,7 +201,7 @@ class TorrentManager:
     ):
         loop = asyncio.get_running_loop()
         self.ranked_torrents = await loop.run_in_executor(
-            None,
+            get_executor(),
             rank_worker,
             self.torrents,
             self.debrid_service,

comet/services/torrent_manager.py

@@ -1,11 +1,10 @@
 import asyncio
 import base64
 import hashlib
-import html
 import re
 import time
 from collections import defaultdict
-from urllib.parse import parse_qs, urlparse
+from urllib.parse import unquote
 
 import aiohttp
 import anyio
@@ -20,15 +19,14 @@ from comet.core.logger import logger
 from comet.core.models import database, settings
 from comet.utils.parsing import default_dump, is_video
 
+TRACKER_PATTERN = re.compile(r"[&?]tr=([^&]+)")
 INFO_HASH_PATTERN = re.compile(r"btih:([a-fA-F0-9]{40}|[a-zA-Z0-9]{32})")
 
 
 def extract_trackers_from_magnet(magnet_uri: str):
     try:
-        decoded_uri = html.unescape(magnet_uri)
-        parsed = urlparse(decoded_uri)
-        params = parse_qs(parsed.query)
-        return params.get("tr", [])
+        trackers = TRACKER_PATTERN.findall(magnet_uri)
+        return [unquote(tracker) for tracker in trackers]
     except Exception as e:
         logger.warning(f"Failed to extract trackers from magnet URI: {e}")
         return []
@@ -247,7 +245,7 @@ add_torrent_queue = AddTorrentQueue()
 
 
 class TorrentUpdateQueue:
-    def __init__(self, batch_size: int = 100, flush_interval: float = 5.0):
+    def __init__(self, batch_size: int = 1000, flush_interval: float = 5.0):
         self.queue = asyncio.Queue()
         self.batch_size = batch_size
         self.flush_interval = flush_interval