diff --git a/comet/api/app.py b/comet/api/app.py index 1f20dc5..1e885c6 100644 --- a/comet/api/app.py +++ b/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 = [ diff --git a/comet/api/endpoints/stream.py b/comet/api/endpoints/stream.py index d22b5e7..c464e72 100644 --- a/comet/api/endpoints/stream.py +++ b/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( diff --git a/comet/core/execution.py b/comet/core/execution.py new file mode 100644 index 0000000..8967430 --- /dev/null +++ b/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 diff --git a/comet/debrid/exceptions.py b/comet/debrid/exceptions.py new file mode 100644 index 0000000..867c553 --- /dev/null +++ b/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) diff --git a/comet/debrid/stremthru.py b/comet/debrid/stremthru.py index 43a796d..561fac2 100644 --- a/comet/debrid/stremthru.py +++ b/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] diff --git a/comet/main.py b/comet/main.py index bac6f53..260e718 100644 --- a/comet/main.py +++ b/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(): diff --git a/comet/services/orchestration.py b/comet/services/orchestration.py index ba57094..bc0c1ec 100644 --- a/comet/services/orchestration.py +++ b/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, diff --git a/comet/services/torrent_manager.py b/comet/services/torrent_manager.py index c180973..25f4e3f 100644 --- a/comet/services/torrent_manager.py +++ b/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