diff --git a/comet/api/app.py b/comet/api/app.py index 5494d98..f7ebd48 100644 --- a/comet/api/app.py +++ b/comet/api/app.py @@ -19,6 +19,8 @@ from comet.core.logger import logger from comet.core.models import settings from comet.services.anime import anime_mapper from comet.services.bandwidth import bandwidth_monitor +from comet.services.torrent_manager import (add_torrent_queue, + torrent_update_queue) from comet.services.trackers import download_best_trackers @@ -89,6 +91,10 @@ async def lifespan(app: FastAPI): if settings.PROXY_DEBRID_STREAM: await bandwidth_monitor.shutdown() + await anime_mapper.stop() + await add_torrent_queue.stop() + await torrent_update_queue.stop() + await teardown_database() diff --git a/comet/core/database.py b/comet/core/database.py index b2bb828..410b1ef 100644 --- a/comet/core/database.py +++ b/comet/core/database.py @@ -647,7 +647,11 @@ async def _run_startup_cleanup(): return current_time = time.time() - should_run = True if interval == 0 else await _should_run_startup_cleanup(current_time, interval) + should_run = ( + True + if interval == 0 + else await _should_run_startup_cleanup(current_time, interval) + ) if not should_run: logger.log("DATABASE", "Startup cleanup skipped (recent run)") return diff --git a/comet/core/models.py b/comet/core/models.py index b548ee0..e766c71 100644 --- a/comet/core/models.py +++ b/comet/core/models.py @@ -13,6 +13,7 @@ from RTN.models import (AudioRankModel, CustomRank, CustomRanksConfig, RipsRankModel) from comet.core.db_router import ReplicaAwareDatabase +from comet.core.logger import logger class AppSettings(BaseSettings): @@ -726,9 +727,7 @@ if settings.DATABASE_TYPE != "sqlite" and settings.DATABASE_READ_REPLICA_URLS: if replica_url: replica_instances.append(_build_database_instance(replica_url)) elif settings.DATABASE_TYPE == "sqlite" and settings.DATABASE_READ_REPLICA_URLS: - logger.log( - "DATABASE", "Read replicas are ignored for sqlite deployments" - ) + logger.log("DATABASE", "Read replicas are ignored for sqlite deployments") database = ReplicaAwareDatabase( _build_database_instance(database_url), replicas=replica_instances diff --git a/comet/main.py b/comet/main.py index e0ac51a..bac6f53 100644 --- a/comet/main.py +++ b/comet/main.py @@ -1,6 +1,5 @@ import contextlib import os -import signal import sys import threading import time @@ -34,17 +33,6 @@ class Server(uvicorn.Server): sys.exit(0) -def signal_handler(sig, frame): - # This will handle kubernetes/docker shutdowns better - # Toss anything that needs to be gracefully shutdown here - logger.log("COMET", "Exiting Gracefully.") - sys.exit(0) - - -signal.signal(signal.SIGINT, signal_handler) -signal.signal(signal.SIGTERM, signal_handler) - - def run_with_uvicorn(): """Run the server with uvicorn only""" config = uvicorn.Config( diff --git a/comet/services/anime.py b/comet/services/anime.py index 8362668..9fc3b35 100644 --- a/comet/services/anime.py +++ b/comet/services/anime.py @@ -111,9 +111,7 @@ class AnimeMapper: if self._background_task and not self._background_task.done(): return - self._background_task = asyncio.create_task( - self._refresh_loop(interval) - ) + self._background_task = asyncio.create_task(self._refresh_loop(interval)) async def _refresh_from_remote( self, @@ -242,5 +240,13 @@ class AnimeMapper: except Exception as exc: logger.warning(f"Anime mapping refresh loop encountered an error: {exc}") + async def stop(self): + if self._background_task: + self._background_task.cancel() + try: + await self._background_task + except asyncio.CancelledError: + pass + anime_mapper = AnimeMapper() diff --git a/comet/services/torrent_manager.py b/comet/services/torrent_manager.py index 0602152..ed06c42 100644 --- a/comet/services/torrent_manager.py +++ b/comet/services/torrent_manager.py @@ -4,6 +4,7 @@ import hashlib import html import re import time +from collections import defaultdict from urllib.parse import parse_qs, urlparse import aiohttp @@ -14,8 +15,6 @@ from demagnetize.core import Demagnetizer from RTN import ParsedData, parse from torf import Magnet -from collections import defaultdict - from comet.core.logger import logger from comet.core.models import database, settings from comet.utils.parsing import default_dump, is_video @@ -237,6 +236,11 @@ class AddTorrentQueue: self.is_running = False + async def stop(self): + self.is_running = False + if not self.queue.empty(): + await self.queue.join() + add_torrent_queue = AddTorrentQueue() @@ -295,6 +299,21 @@ class TorrentUpdateQueue: self.is_running = False + async def stop(self): + self.is_running = False + + # Process remaining items in queue + while not self.queue.empty(): + try: + file_info, media_id = self.queue.get_nowait() + await self._process_file_info(file_info, media_id) + except Exception: + break + + # Flush any remaining batches + if any(len(batch) > 0 for batch in self.batches.values()): + await self._flush_batch() + def _reset_batches(self): for key in self.batches: if len(self.batches[key]) > 0: