From 04dfd38c8f892af72d33641ab3b8af900e68b3e3 Mon Sep 17 00:00:00 2001 From: g0ldyy <153996346+g0ldyy@users.noreply.github.com> Date: Mon, 5 Jan 2026 14:50:28 +0100 Subject: [PATCH] feat: switch from session-level to transaction-level PostgreSQL advisory locks for database cleanup and batched upserts --- comet/core/database.py | 140 ++++++++++++++---------------- comet/services/torrent_manager.py | 58 +++++-------- 2 files changed, 91 insertions(+), 107 deletions(-) diff --git a/comet/core/database.py b/comet/core/database.py index 83fdbaa..2466432 100644 --- a/comet/core/database.py +++ b/comet/core/database.py @@ -560,69 +560,86 @@ async def _run_startup_cleanup(): logger.log("DATABASE", "Startup cleanup skipped (recent run)") return - lock_acquired = await _try_acquire_startup_cleanup_lock() - if not lock_acquired: - logger.log("DATABASE", "Startup cleanup already running elsewhere; skipping") - return - try: - logger.log("DATABASE", "Running startup cleanup sweep") + async with database.transaction(): + if settings.DATABASE_TYPE == "postgresql": + acquired = await database.fetch_val( + "SELECT pg_try_advisory_xact_lock(:lock_id)", + {"lock_id": STARTUP_CLEANUP_LOCK_ID}, + force_primary=True, + ) + if not acquired: + logger.log( + "DATABASE", + "Startup cleanup already running elsewhere; skipping", + ) + return - await database.execute( - """ - DELETE FROM first_searches - WHERE timestamp < CAST(:current_time AS BIGINT) - CAST(:cache_ttl AS BIGINT); - """, - {"cache_ttl": settings.TORRENT_CACHE_TTL, "current_time": current_time}, - ) + logger.log("DATABASE", "Running startup cleanup sweep") - await database.execute( - """ - DELETE FROM metadata_cache - WHERE timestamp < CAST(:current_time AS BIGINT) - CAST(:cache_ttl AS BIGINT); - """, - {"cache_ttl": settings.METADATA_CACHE_TTL, "current_time": current_time}, - ) - - if settings.TORRENT_CACHE_TTL >= 0: await database.execute( """ - DELETE FROM torrents + DELETE FROM first_searches WHERE timestamp < CAST(:current_time AS BIGINT) - CAST(:cache_ttl AS BIGINT); """, {"cache_ttl": settings.TORRENT_CACHE_TTL, "current_time": current_time}, ) - await database.execute( - """ - DELETE FROM debrid_availability - WHERE timestamp < CAST(:current_time AS BIGINT) - CAST(:cache_ttl AS BIGINT); - """, - {"cache_ttl": settings.DEBRID_CACHE_TTL, "current_time": current_time}, - ) + await database.execute( + """ + DELETE FROM metadata_cache + WHERE timestamp < CAST(:current_time AS BIGINT) - CAST(:cache_ttl AS BIGINT); + """, + { + "cache_ttl": settings.METADATA_CACHE_TTL, + "current_time": current_time, + }, + ) - await database.execute( - """ - DELETE FROM digital_release_cache - WHERE timestamp < CAST(:current_time AS BIGINT) - CAST(:cache_ttl AS BIGINT); - """, - {"cache_ttl": settings.METADATA_CACHE_TTL, "current_time": current_time}, - ) + if settings.TORRENT_CACHE_TTL >= 0: + await database.execute( + """ + DELETE FROM torrents + WHERE timestamp < CAST(:current_time AS BIGINT) - CAST(:cache_ttl AS BIGINT); + """, + { + "cache_ttl": settings.TORRENT_CACHE_TTL, + "current_time": current_time, + }, + ) - await database.execute("DELETE FROM download_links_cache") + await database.execute( + """ + DELETE FROM debrid_availability + WHERE timestamp < CAST(:current_time AS BIGINT) - CAST(:cache_ttl AS BIGINT); + """, + {"cache_ttl": settings.DEBRID_CACHE_TTL, "current_time": current_time}, + ) - await database.execute( - """ - INSERT INTO db_maintenance (id, last_startup_cleanup) - VALUES (1, :timestamp) - ON CONFLICT (id) DO UPDATE SET last_startup_cleanup = :timestamp - """, - {"timestamp": current_time}, - force_primary=True, - ) - finally: - if lock_acquired: - await _release_startup_cleanup_lock() + await database.execute( + """ + DELETE FROM digital_release_cache + WHERE timestamp < CAST(:current_time AS BIGINT) - CAST(:cache_ttl AS BIGINT); + """, + { + "cache_ttl": settings.METADATA_CACHE_TTL, + "current_time": current_time, + }, + ) + + await database.execute("DELETE FROM download_links_cache") + + await database.execute( + """ + INSERT INTO db_maintenance (id, last_startup_cleanup) + VALUES (1, :timestamp) + ON CONFLICT (id) DO UPDATE SET last_startup_cleanup = :timestamp + """, + {"timestamp": current_time}, + force_primary=True, + ) + except Exception as e: + logger.error(f"Error executing startup cleanup: {e}") async def _should_run_startup_cleanup(current_time: float, interval: int) -> bool: @@ -637,29 +654,6 @@ async def _should_run_startup_cleanup(current_time: float, interval: int) -> boo return (current_time - last_run) >= interval -async def _try_acquire_startup_cleanup_lock() -> bool: - if settings.DATABASE_TYPE != "postgresql": - return True - - result = await database.fetch_val( - "SELECT pg_try_advisory_lock(:lock_id)", - {"lock_id": STARTUP_CLEANUP_LOCK_ID}, - force_primary=True, - ) - return bool(result) - - -async def _release_startup_cleanup_lock(): - if settings.DATABASE_TYPE != "postgresql": - return - - await database.fetch_val( - "SELECT pg_advisory_unlock(:lock_id)", - {"lock_id": STARTUP_CLEANUP_LOCK_ID}, - force_primary=True, - ) - - async def cleanup_expired_locks(): while True: try: diff --git a/comet/services/torrent_manager.py b/comet/services/torrent_manager.py index 54c7c05..f1a02e5 100644 --- a/comet/services/torrent_manager.py +++ b/comet/services/torrent_manager.py @@ -642,46 +642,36 @@ async def _execute_batched_upsert(query: str, rows): return ordered_rows = sorted(rows, key=lambda row: row.get("lock_key") or 0) - - acquired_locks = [] rows_to_insert = [] 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 + async with database.transaction(): + for row in ordered_rows: + lock_key = row.get("lock_key") + if lock_key is None: + rows_to_insert.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 - - if rows_to_insert: - sanitized_rows = [ - {key: value for key, value in row.items() if key != "lock_key"} - for row in rows_to_insert - ] - - await database.execute_many(query, sanitized_rows) - - 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))", + # Use transaction-level non-blocking lock + # This automatically releases at the end of the transaction + acquired = await database.fetch_val( + "SELECT pg_try_advisory_xact_lock(CAST(:lock_key AS BIGINT))", {"lock_key": lock_key}, ) - except Exception: - pass # Best effort unlock + if acquired: + rows_to_insert.append(row) + # If not acquired, skip this row - another replica is handling it + + if rows_to_insert: + sanitized_rows = [ + {key: value for key, value in row.items() if key != "lock_key"} + for row in rows_to_insert + ] + + await database.execute_many(query, sanitized_rows) + + except Exception as e: + logger.warning(f"Error executing batched upsert: {e}") def _get_torrent_upsert_query(conflict_key: str) -> str: