refactor: refactor shutdown handling and improve cleanup processes

- Removed signal handling from main.py for graceful shutdown.
- Added stop methods to AnimeMapper, AddTorrentQueue, and TorrentUpdateQueue for better resource management during shutdown.
- Updated lifespan in app.py to ensure proper shutdown of services.

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()
 
 

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

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

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(

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()

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: