mirror of
https://github.com/g0ldyy/comet.git
synced 2026-01-12 01:16:12 +01:00
feat: switch from session-level to transaction-level PostgreSQL advisory locks for database cleanup and batched upserts
This commit is contained in:
g0ldyy
+67
-73
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user