diff --git a/comet/core/database.py b/comet/core/database.py index e8a1de3..83fdbaa 100644 --- a/comet/core/database.py +++ b/comet/core/database.py @@ -88,15 +88,6 @@ async def setup_database(): """ ) - await database.execute( - """ - CREATE TABLE IF NOT EXISTS ongoing_searches ( - media_id TEXT PRIMARY KEY, - timestamp INTEGER - ) - """ - ) - await database.execute( """ CREATE TABLE IF NOT EXISTS scrape_locks ( @@ -521,6 +512,16 @@ async def setup_database(): """ ) + if settings.DATABASE_TYPE == "postgresql": + # Covering index for get_cached_torrents + await database.execute( + """ + CREATE INDEX IF NOT EXISTS idx_torrents_covering + ON torrents (media_id, season, episode) + INCLUDE (info_hash, file_index, title, seeders, size, tracker, sources, parsed, timestamp) + """ + ) + if settings.DATABASE_TYPE == "sqlite": await database.execute("PRAGMA busy_timeout=30000") # 30 seconds timeout await database.execute("PRAGMA journal_mode=WAL") @@ -534,7 +535,6 @@ async def setup_database(): await database.execute("PRAGMA secure_delete=OFF") await database.execute("PRAGMA auto_vacuum=OFF") - await database.execute("DELETE FROM ongoing_searches") await database.execute("DELETE FROM active_connections") await database.execute("DELETE FROM metrics_cache") diff --git a/comet/services/debrid_cache.py b/comet/services/debrid_cache.py index 28efb71..34fa235 100644 --- a/comet/services/debrid_cache.py +++ b/comet/services/debrid_cache.py @@ -5,6 +5,25 @@ import orjson from comet.core.models import database, settings from comet.utils.parsing import default_dump +DEBRID_UPDATE_INTERVAL = ( + settings.DEBRID_CACHE_TTL // 2 if settings.DEBRID_CACHE_TTL > 0 else 31536000 +) + +CONDITIONAL_UPDATE = """ + DO UPDATE SET + title = EXCLUDED.title, + file_index = EXCLUDED.file_index, + size = EXCLUDED.size, + parsed = EXCLUDED.parsed, + timestamp = EXCLUDED.timestamp + WHERE + debrid_availability.title IS DISTINCT FROM EXCLUDED.title + OR debrid_availability.file_index IS DISTINCT FROM EXCLUDED.file_index + OR debrid_availability.size IS DISTINCT FROM EXCLUDED.size + OR debrid_availability.parsed IS DISTINCT FROM EXCLUDED.parsed + OR COALESCE(debrid_availability.timestamp, 0) < (EXCLUDED.timestamp - :update_interval) +""" + async def cache_availability(debrid_service: str, availability: list): current_time = time.time() @@ -22,6 +41,7 @@ async def cache_availability(debrid_service: str, availability: list): if file["parsed"] is not None else None, "timestamp": current_time, + "update_interval": DEBRID_UPDATE_INTERVAL, } for file in availability ] @@ -32,8 +52,11 @@ async def cache_availability(debrid_service: str, availability: list): INTO debrid_availability VALUES (:debrid_service, :info_hash, :file_index, :title, :season, :episode, :size, :parsed, :timestamp) """ - await database.execute_many(query, values) - elif settings.DATABASE_TYPE == "postgresql": + sqlite_values = [ + {k: v for k, v in val.items() if k != "update_interval"} for val in values + ] + await database.execute_many(query, sqlite_values) + else: both_values = [] season_only_values = [] episode_only_values = [] @@ -49,73 +72,45 @@ async def cache_availability(debrid_service: str, availability: list): else: no_season_episode_values.append(val) - # handle each case separately with appropriate ON CONFLICT clauses if both_values: - query = """ + query = f""" INSERT INTO debrid_availability VALUES (:debrid_service, :info_hash, :file_index, :title, :season, :episode, :size, :parsed, :timestamp) ON CONFLICT (debrid_service, info_hash, season, episode) WHERE season IS NOT NULL AND episode IS NOT NULL - DO UPDATE SET - title = EXCLUDED.title, - file_index = EXCLUDED.file_index, - size = EXCLUDED.size, - parsed = EXCLUDED.parsed, - timestamp = EXCLUDED.timestamp + {CONDITIONAL_UPDATE} """ await database.execute_many(query, both_values) if season_only_values: - query = """ + query = f""" INSERT INTO debrid_availability VALUES (:debrid_service, :info_hash, :file_index, :title, :season, :episode, :size, :parsed, :timestamp) ON CONFLICT (debrid_service, info_hash, season) WHERE season IS NOT NULL AND episode IS NULL - DO UPDATE SET - title = EXCLUDED.title, - file_index = EXCLUDED.file_index, - size = EXCLUDED.size, - parsed = EXCLUDED.parsed, - timestamp = EXCLUDED.timestamp + {CONDITIONAL_UPDATE} """ await database.execute_many(query, season_only_values) if episode_only_values: - query = """ + query = f""" INSERT INTO debrid_availability VALUES (:debrid_service, :info_hash, :file_index, :title, :season, :episode, :size, :parsed, :timestamp) ON CONFLICT (debrid_service, info_hash, episode) WHERE season IS NULL AND episode IS NOT NULL - DO UPDATE SET - title = EXCLUDED.title, - file_index = EXCLUDED.file_index, - size = EXCLUDED.size, - parsed = EXCLUDED.parsed, - timestamp = EXCLUDED.timestamp + {CONDITIONAL_UPDATE} """ await database.execute_many(query, episode_only_values) if no_season_episode_values: - query = """ + query = f""" INSERT INTO debrid_availability VALUES (:debrid_service, :info_hash, :file_index, :title, :season, :episode, :size, :parsed, :timestamp) ON CONFLICT (debrid_service, info_hash) WHERE season IS NULL AND episode IS NULL - DO UPDATE SET - title = EXCLUDED.title, - file_index = EXCLUDED.file_index, - size = EXCLUDED.size, - parsed = EXCLUDED.parsed, - timestamp = EXCLUDED.timestamp + {CONDITIONAL_UPDATE} """ await database.execute_many(query, no_season_episode_values) - else: - query = """ - INSERT - INTO debrid_availability - VALUES (:debrid_service, :info_hash, :file_index, :title, :season, :episode, :size, :parsed, :timestamp) - """ - await database.execute_many(query, values) async def get_cached_availability( diff --git a/comet/services/torrent_manager.py b/comet/services/torrent_manager.py index 7904944..d57ea70 100644 --- a/comet/services/torrent_manager.py +++ b/comet/services/torrent_manager.py @@ -243,6 +243,11 @@ class AddTorrentQueue: add_torrent_queue = AddTorrentQueue() +UPDATE_INTERVAL = ( + settings.TORRENT_CACHE_TTL // 2 if settings.TORRENT_CACHE_TTL >= 0 else 31536000 +) + + class TorrentUpdateQueue: def __init__(self, batch_size: int = 1000, flush_interval: float = 5.0): self.queue = asyncio.Queue() @@ -641,74 +646,61 @@ async def _execute_batched_upsert(query: str, rows): await _execute_sqlite_batched_upsert(rows) return - ordered_rows = sorted(rows, key=lambda row: row.get("lock_key")) + ordered_rows = sorted(rows, key=lambda row: row.get("lock_key") or 0) - attempts = ( - POSTGRES_LOCK_RETRY_ATTEMPTS + 1 - if settings.DATABASE_TYPE == "postgresql" - else 1 - ) + acquired_locks = [] + rows_to_insert = [] - for attempt in range(attempts): - try: - async with database.transaction(): - rows_to_insert = ordered_rows + try: + # Non-blocking lock acquisition - skip rows we can't lock + for row in ordered_rows: + lock_key = row.get("lock_key") + if lock_key is None: + rows_to_insert.append(row) + continue - if settings.DATABASE_TYPE == "postgresql": - locked_rows = [] - for row in ordered_rows: - lock_key = row.get("lock_key") - if lock_key is None: - locked_rows.append(row) - continue + # Use session-level non-blocking lock (not transaction-level) + acquired = await database.fetch_val( + "SELECT pg_try_advisory_lock(CAST(:lock_key AS BIGINT))", + {"lock_key": lock_key}, + ) + if acquired: + acquired_locks.append(lock_key) + rows_to_insert.append(row) + # If not acquired, skip this row - another replica is handling it - acquired = await database.fetch_val( - "SELECT pg_try_advisory_xact_lock(CAST(:lock_key AS BIGINT))", - {"lock_key": lock_key}, - ) - if acquired: - locked_rows.append(row) + if rows_to_insert: + sanitized_rows = [ + {key: value for key, value in row.items() if key != "lock_key"} + for row in rows_to_insert + ] - rows_to_insert = locked_rows + if sanitized_rows: + # No retry loop needed - we only process rows we have locks for + await database.execute_many(query, sanitized_rows) - sanitized_rows = [ - {key: value for key, value in row.items() if key != "lock_key"} - for row in rows_to_insert - ] - - if sanitized_rows: - 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)) + finally: + # Always release all acquired session-level locks + for lock_key in acquired_locks: + try: + await database.execute( + "SELECT pg_advisory_unlock(CAST(:lock_key AS BIGINT))", + {"lock_key": lock_key}, + ) + except Exception: + pass # Best effort unlock def _get_torrent_upsert_query(conflict_key: str) -> str: if settings.DATABASE_TYPE == "sqlite": return SQLITE_UPSERT_QUERY - if settings.DATABASE_TYPE == "postgresql": - target = POSTGRES_CONFLICT_TARGETS[conflict_key] - if conflict_key not in _POSTGRES_UPSERT_CACHE: - _POSTGRES_UPSERT_CACHE[conflict_key] = ( - TORRENT_INSERT_TEMPLATE - + f" ON CONFLICT {target} " - + POSTGRES_UPDATE_SET - ) - return _POSTGRES_UPSERT_CACHE[conflict_key] - - return TORRENT_INSERT_TEMPLATE - - -UPDATE_INTERVAL = 31536000 # Default 1 year -if settings.LIVE_TORRENT_CACHE_TTL >= 0: - UPDATE_INTERVAL = settings.LIVE_TORRENT_CACHE_TTL // 2 + target = POSTGRES_CONFLICT_TARGETS[conflict_key] + if conflict_key not in _POSTGRES_UPSERT_CACHE: + _POSTGRES_UPSERT_CACHE[conflict_key] = ( + TORRENT_INSERT_TEMPLATE + f" ON CONFLICT {target} " + POSTGRES_UPDATE_SET + ) + return _POSTGRES_UPSERT_CACHE[conflict_key] async def _upsert_torrent_record(params: dict): @@ -722,21 +714,7 @@ async def _upsert_torrent_record(params: dict): params["update_interval"] = UPDATE_INTERVAL - 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)) + await database.execute(query, params) def _is_retryable_lock_error(exc: Exception) -> bool: