g0ldyy/comet
1
0
Fork 0
mirror of https://github.com/g0ldyy/comet.git synced 2026-01-12 01:16:12 +01:00
Code Issues Packages Projects Releases Wiki Activity

feat: improve torrent batch processing with in-memory deduplication, PostgreSQL advisory locks, and enhance metadata handling

Browse Source
This commit is contained in:
g0ldyy
2025-12-31 13:27:48 +01:00
parent 356be2590c
commit 45ff7c173a
7 changed files with 189 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comet/api/endpoints/config.py
+1
View File
@@ -29,5 +29,6 @@ async def configure(request: Request):
else "",
"webConfig": web_config,
"proxyDebridStream": settings.PROXY_DEBRID_STREAM,
"disableTorrentStreams": settings.DISABLE_TORRENT_STREAMS,
},
)
comet/api/endpoints/playback.py
+3 -2
View File
@@ -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(
comet/core/database.py
+3 -1
View File
@@ -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}")
comet/metadata/manager.py
+3
View File
@@ -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
comet/services/torrent_manager.py
+126 -20
View File
@@ -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()
comet/templates/index.html
+39 -15
View File
@@ -409,16 +409,26 @@
help-text="Debrid Stream Proxying allows you to use your Debrid Service from multiple IPs at same time!"></sl-input>
</div>
{% set default_debrid_service = 'realdebrid' if disableTorrentStreams else 'torrent' %}
<div class="form-item">
<sl-select id="debridService" value="torrent" label="Debrid Service" placeholder="Select debrid service">
<sl-select
id="debridService"
value="{{ default_debrid_service }}"
data-default-service="{{ default_debrid_service }}"
data-disable-torrent="{{ 'true' if disableTorrentStreams else 'false' }}"
label="Debrid Service"
placeholder="Select debrid service"
>
{% if not disableTorrentStreams %}
<sl-option value="torrent">Torrent</sl-option>
{% endif %}
<sl-option value="realdebrid">Real-Debrid</sl-option>
<sl-option value="torbox">TorBox</sl-option>
<sl-option value="alldebrid">All-Debrid</sl-option>
<sl-option value="debridlink">Debrid-Link</sl-option>
<sl-option value="premiumize">Premiumize</sl-option>
<sl-option value="debrider">Debrider</sl-option>
<sl-option value="easydebrid">EasyDebrid</sl-option>
<sl-option value="realdebrid">Real-Debrid</sl-option>
<sl-option value="debridlink">Debrid-Link</sl-option>
<sl-option value="alldebrid">All-Debrid</sl-option>
<sl-option value="premiumize">Premiumize</sl-option>
<sl-option value="offcloud">Offcloud</sl-option>
<sl-option value="pikpak">PikPak</sl-option>
</sl-select>
@@ -455,9 +465,13 @@
</sl-details>
<script>
document
.getElementById("debridService")
.addEventListener("sl-change", function (event) {
const debridServiceElement = document.getElementById("debridService");
const torrentDisabled = debridServiceElement.dataset.disableTorrent === "true";
const defaultDebridService = debridServiceElement.dataset.defaultService;
window.disableTorrentStreams = torrentDisabled;
window.defaultDebridService = defaultDebridService;
debridServiceElement.addEventListener("sl-change", function (event) {
const selectedService = event.target.value;
const apiKeyLink = document.getElementById("apiKeyLink");
const apiKeyInput = document.getElementById("debridApiKey");
@@ -528,13 +542,13 @@
selectedService === "pikpak"
) {
apiKeyInput.helpText = "Format: `email:password`";
} else if (selectedService != "torrent") {
apiKeyInput.helpText = "Format: `api-key`";
} else {
} else if (!torrentDisabled && selectedService === "torrent") {
apiKeyInput.helpText = "";
} else {
apiKeyInput.helpText = "Format: `api-key`";
}
if (selectedService === "torrent") {
if (!torrentDisabled && selectedService === "torrent") {
apiKeyInput.disabled = true;
} else {
apiKeyInput.disabled = false;
@@ -746,7 +760,10 @@
const allowEnglishInLanguages = document.getElementById("allowEnglishInLanguages").checked;
const removeUnknownLanguages = document.getElementById("removeUnknownLanguages").checked;
const resultFormat = Array.from(document.getElementById("resultFormat").selectedOptions).map(option => option.value);
const debridService = document.getElementById("debridService").value;
let debridService = document.getElementById("debridService").value;
if (window.disableTorrentStreams && debridService === "torrent") {
debridService = window.defaultDebridService;
}
const debridApiKey = document.getElementById("debridApiKey").value;
const debridStreamProxyPassword = document.getElementById("debridStreamProxyPassword").value;
@@ -824,8 +841,15 @@
document.getElementById("maxResultsPerResolution").value = settings.maxResultsPerResolution;
if (settings.maxSize !== null)
document.getElementById("maxSize").value = settings.maxSize / 1073741824;
if (settings.debridService !== null)
document.getElementById("debridService").value = settings.debridService;
if (settings.debridService !== null) {
const debridServiceSelect = document.getElementById("debridService");
const requestedService = settings.debridService;
if (window.disableTorrentStreams && requestedService === "torrent") {
debridServiceSelect.value = window.defaultDebridService;
} else {
debridServiceSelect.value = requestedService;
}
}
if (settings.debridApiKey !== null)
document.getElementById("debridApiKey").value = settings.debridApiKey;
if (settings.debridStreamProxyPassword !== null)
comet/utils/parsing.py
+14 -2
View File
@@ -49,18 +49,30 @@ def default_dump(obj):
return obj.model_dump()
def parse_optional_int(value: str | None):
if value == "n" or value is None or value == "":
return None
try:
return int(value)
except ValueError:
return None
def parse_media_id(media_type: str, media_id: str):
if "kitsu" in media_id:
info = media_id.split(":")
if len(info) > 2:
return info[1], 1, int(info[2])
return info[1], 1, parse_optional_int(info[2])
else:
return info[1], 1, None
if media_type == "series":
info = media_id.split(":")
return info[0], int(info[1]), int(info[2])
series_id = info[0]
season = parse_optional_int(info[1]) if len(info) > 1 else None
episode = parse_optional_int(info[2]) if len(info) > 2 else None
return series_id, season, episode
return media_id, None, None