From 45ff7c173abb49a050c4ba14103fe6beae845d32 Mon Sep 17 00:00:00 2001 From: g0ldyy <153996346+g0ldyy@users.noreply.github.com> Date: Wed, 31 Dec 2025 13:27:48 +0100 Subject: [PATCH] feat: improve torrent batch processing with in-memory deduplication, PostgreSQL advisory locks, and enhance metadata handling --- comet/api/endpoints/config.py | 1 + comet/api/endpoints/playback.py | 5 +- comet/core/database.py | 4 +- comet/metadata/manager.py | 3 + comet/services/torrent_manager.py | 146 ++++++++++++++++++++++++++---- comet/templates/index.html | 54 ++++++++--- comet/utils/parsing.py | 16 +++- 7 files changed, 189 insertions(+), 40 deletions(-) diff --git a/comet/api/endpoints/config.py b/comet/api/endpoints/config.py index 5948b0c..97525de 100644 --- a/comet/api/endpoints/config.py +++ b/comet/api/endpoints/config.py @@ -29,5 +29,6 @@ async def configure(request: Request): else "", "webConfig": web_config, "proxyDebridStream": settings.PROXY_DEBRID_STREAM, + "disableTorrentStreams": settings.DISABLE_TORRENT_STREAMS, }, ) diff --git a/comet/api/endpoints/playback.py b/comet/api/endpoints/playback.py index eef73b8..020b031 100644 --- a/comet/api/endpoints/playback.py +++ b/comet/api/endpoints/playback.py @@ -12,6 +12,7 @@ from comet.debrid.manager import get_debrid from comet.metadata.manager import MetadataScraper from comet.services.streaming.manager import custom_handle_stream_request from comet.utils.network import NO_CACHE_HEADERS, get_client_ip +from comet.utils.parsing import parse_optional_int router = APIRouter() @@ -34,8 +35,8 @@ async def playback( ): config = config_check(b64config) - season = int(season) if season != "n" else None - episode = int(episode) if episode != "n" else None + season = parse_optional_int(season) + episode = parse_optional_int(episode) async with aiohttp.ClientSession() as session: cached_link = await database.fetch_one( diff --git a/comet/core/database.py b/comet/core/database.py index 51811ad..a9eb5d1 100644 --- a/comet/core/database.py +++ b/comet/core/database.py @@ -772,7 +772,9 @@ async def _migrate_indexes(): for index_name in old_indexes: await database.execute(f"DROP INDEX IF EXISTS {index_name}") - logger.log("COMET", "Database: Legacy indexes dropped. New indexes will be created.") + logger.log( + "COMET", "Database: Legacy indexes dropped. New indexes will be created." + ) except Exception as e: logger.warning(f"Error during index migration: {e}") diff --git a/comet/metadata/manager.py b/comet/metadata/manager.py index 46c15a7..c9ba36d 100644 --- a/comet/metadata/manager.py +++ b/comet/metadata/manager.py @@ -83,6 +83,9 @@ class MetadataScraper: ) def normalize_metadata(self, metadata: dict, season: int, episode: int): + if not metadata: + return None + title, year, year_end = metadata if title is None: # metadata retrieving failed diff --git a/comet/services/torrent_manager.py b/comet/services/torrent_manager.py index 25f4e3f..441d95d 100644 --- a/comet/services/torrent_manager.py +++ b/comet/services/torrent_manager.py @@ -250,7 +250,7 @@ class TorrentUpdateQueue: self.batch_size = batch_size self.flush_interval = flush_interval self.is_running = False - self.batches = {"to_delete": [], "upserts": []} + self.batches = {"to_delete": set(), "upserts": {}} async def add_torrent_info(self, file_info: dict, media_id: str = None): await self.queue.put((file_info, media_id)) @@ -317,30 +317,32 @@ class TorrentUpdateQueue: await self._flush_batch() def _reset_batches(self): - for key in self.batches: - if len(self.batches[key]) > 0: + for key, batch in self.batches.items(): + if len(batch) > 0: logger.warning( - f"Ignoring {len(self.batches[key])} items in problematic '{key}' batch" + f"Ignoring {len(batch)} items in problematic '{key}' batch" ) - self.batches[key] = [] + batch.clear() async def _flush_batch(self): try: if self.batches["to_delete"]: + delete_items = list(self.batches["to_delete"]) sub_batch_size = 100 - for i in range(0, len(self.batches["to_delete"]), sub_batch_size): + for i in range(0, len(delete_items), sub_batch_size): try: - sub_batch = self.batches["to_delete"][i : i + sub_batch_size] + sub_batch = delete_items[i : i + sub_batch_size] placeholders = [] params = {} for idx, item in enumerate(sub_batch): + info_hash, season = item key_suffix = f"_{idx}" placeholders.append( f"(CAST(:info_hash{key_suffix} AS TEXT), CAST(:season{key_suffix} AS INTEGER))" ) - params[f"info_hash{key_suffix}"] = item["info_hash"] - params[f"season{key_suffix}"] = item["season"] + params[f"info_hash{key_suffix}"] = info_hash + params[f"season{key_suffix}"] = season async with database.transaction(): delete_query = f""" @@ -354,28 +356,28 @@ class TorrentUpdateQueue: except Exception as e: logger.warning(f"Error processing delete batch: {e}") - self.batches["to_delete"] = [] + self.batches["to_delete"].clear() if self.batches["upserts"]: grouped: dict[str, list[dict]] = defaultdict(list) - for params in self.batches["upserts"]: + for params in self.batches["upserts"].values(): key = _determine_conflict_key(params["season"], params["episode"]) grouped[key].append(params) for key, rows in grouped.items(): query = _get_torrent_upsert_query(key) try: - async with database.transaction(): - await database.execute_many(query, rows) + await _execute_batched_upsert(query, rows) except Exception as e: logger.warning(f"Error processing upsert batch: {e}") - if len(self.batches["upserts"]) > 0: + total_upserts = len(self.batches["upserts"]) + if total_upserts > 0: logger.log( "SCRAPER", - f"Upserted {len(self.batches['upserts'])} torrents in batch", + f"Upserted {total_upserts} torrents in batch", ) - self.batches["upserts"] = [] + self.batches["upserts"].clear() except Exception as e: logger.warning(f"Error in flush_batch: {e}") @@ -400,11 +402,28 @@ class TorrentUpdateQueue: "media_id": media_id, } - self.batches["upserts"].append(params) + params["lock_key"] = _compute_advisory_lock_key( + media_id, + file_info["info_hash"], + file_info["season"], + file_info["episode"], + ) + + upsert_key = _build_upsert_key( + file_info["info_hash"], + file_info["season"], + file_info["episode"], + media_id, + ) + + # In-memory deduplication: keep the freshest timestamp + existing = self.batches["upserts"].get(upsert_key) + if not existing or params["timestamp"] > existing["timestamp"]: + self.batches["upserts"][upsert_key] = params if file_info["episode"] is not None: - self.batches["to_delete"].append( - {"info_hash": file_info["info_hash"], "season": file_info["season"]} + self.batches["to_delete"].add( + (file_info["info_hash"], file_info["season"]) ) await self._check_batch_size() @@ -461,8 +480,13 @@ POSTGRES_UPDATE_SET = """ sources = EXCLUDED.sources, parsed = EXCLUDED.parsed, timestamp = EXCLUDED.timestamp + WHERE COALESCE(torrents.timestamp, 0) < EXCLUDED.timestamp """ +POSTGRES_LOCK_TIMEOUT = "750ms" +POSTGRES_LOCK_RETRY_ATTEMPTS = 1 +POSTGRES_RETRYABLE_SQLSTATES = {"55P03", "40P01"} + POSTGRES_CONFLICT_TARGETS = { "series": "(media_id, info_hash, season, episode) WHERE season IS NOT NULL AND episode IS NOT NULL", "season_only": "(media_id, info_hash, season) WHERE season IS NOT NULL AND episode IS NULL", @@ -483,6 +507,62 @@ def _determine_conflict_key(season, episode) -> str: return "none" +def _build_upsert_key(info_hash, season, episode, media_id): + return (media_id, info_hash, season, episode) + + +def _compute_advisory_lock_key(media_id, info_hash, season, episode) -> int: + payload = f"{media_id}|{info_hash}|{season}|{episode}".encode("utf-8") + digest = hashlib.sha1(payload).digest() + # Use signed=True to get a signed 64-bit int directly, which Postgres expects + return int.from_bytes(digest[:8], byteorder="big", signed=True) + + +async def _execute_batched_upsert(query: str, rows): + if not rows: + return + + ordered_rows = sorted(rows, key=lambda row: row.get("lock_key")) + + sanitized_rows = [ + {key: value for key, value in row.items() if key != "lock_key"} + for row in ordered_rows + ] + + attempts = ( + POSTGRES_LOCK_RETRY_ATTEMPTS + 1 + if settings.DATABASE_TYPE == "postgresql" + else 1 + ) + + for attempt in range(attempts): + try: + async with database.transaction(): + if settings.DATABASE_TYPE == "postgresql": + await database.execute( + f"SET LOCAL lock_timeout = '{POSTGRES_LOCK_TIMEOUT}'" + ) + for row in ordered_rows: + lock_key = row.get("lock_key") + if lock_key is None: + continue + await database.execute( + "SELECT pg_advisory_xact_lock(CAST(:lock_key AS BIGINT))", + {"lock_key": lock_key}, + ) + + await database.execute_many(query, sanitized_rows) + return + except Exception as exc: + if ( + settings.DATABASE_TYPE != "postgresql" + or not _is_retryable_lock_error(exc) + or attempt == attempts - 1 + ): + raise + await asyncio.sleep(0.2 * (attempt + 1)) + + def _get_torrent_upsert_query(conflict_key: str) -> str: if settings.DATABASE_TYPE == "sqlite": return SQLITE_UPSERT_QUERY @@ -504,7 +584,33 @@ async def _upsert_torrent_record(params: dict): query = _get_torrent_upsert_query( _determine_conflict_key(params.get("season"), params.get("episode")) ) - await database.execute(query, params) + if settings.DATABASE_TYPE != "postgresql": + await database.execute(query, params) + return + + for attempt in range(POSTGRES_LOCK_RETRY_ATTEMPTS + 1): + try: + async with database.transaction(): + await database.execute( + f"SET LOCAL lock_timeout = '{POSTGRES_LOCK_TIMEOUT}'" + ) + await database.execute(query, params) + return + except Exception as exc: # pragma: no cover - driver specific + if ( + not _is_retryable_lock_error(exc) + or attempt == POSTGRES_LOCK_RETRY_ATTEMPTS + ): + raise + await asyncio.sleep(0.2 * (attempt + 1)) + + +def _is_retryable_lock_error(exc: Exception) -> bool: + sqlstate = getattr(exc, "sqlstate", None) + if sqlstate in POSTGRES_RETRYABLE_SQLSTATES: + return True + message = str(exc).lower() + return "lock timeout" in message or "deadlock detected" in message torrent_update_queue = TorrentUpdateQueue() diff --git a/comet/templates/index.html b/comet/templates/index.html index cc7dd6c..c1475b8 100644 --- a/comet/templates/index.html +++ b/comet/templates/index.html @@ -409,16 +409,26 @@ help-text="Debrid Stream Proxying allows you to use your Debrid Service from multiple IPs at same time!"> + {% set default_debrid_service = 'realdebrid' if disableTorrentStreams else 'torrent' %}
- + + {% if not disableTorrentStreams %} Torrent + {% endif %} + Real-Debrid TorBox + All-Debrid + Debrid-Link + Premiumize Debrider EasyDebrid - Real-Debrid - Debrid-Link - All-Debrid - Premiumize Offcloud PikPak @@ -455,9 +465,13 @@