From 937bc18ca72bca4fcabd72027d733e0231b4715e Mon Sep 17 00:00:00 2001
From: pacnpal <183241239+pacnpal@users.noreply.github.com>
Date: Tue, 25 Feb 2025 13:09:09 -0500
Subject: [PATCH] Enhance DiscordBot shutdown process; add connection timeout,
 cancel web tasks, and handle SIGINT signal

---
 discord_glhf/bot.py | 33 ++++++++++++++++++++++++---------
 1 file changed, 24 insertions(+), 9 deletions(-)

diff --git a/discord_glhf/bot.py b/discord_glhf/bot.py
index 2bb7c07..96f768e 100644
--- a/discord_glhf/bot.py
+++ b/discord_glhf/bot.py
@@ -124,10 +124,13 @@ class DiscordBot:
         max_retries = 5
         base_delay = 1.0
 
-        while retry_count < max_retries:
+        while retry_count < max_retries and self._running:  # Check if we're still running
             try:
-                await self.bot.connect()
+                # Add a timeout to connection
+                await asyncio.wait_for(self.bot.connect(), timeout=60.0)
                 return
+            except asyncio.TimeoutError:
+                logger.warning("Connection attempt timed out, retrying...")
             except (aiohttp.ClientError, socket.gaierror) as e:
                 retry_count += 1
                 if retry_count == max_retries:
@@ -289,12 +292,22 @@ class DiscordBot:
 
         try:
             async with self._init_lock:
-                # Stop queue processor first
+            # Cancel Hypercorn web tasks first
+                tasks = [t for t in asyncio.all_tasks() if t.get_name() == 'hypercorn_web']
+                for task in tasks:
+                    task.cancel()
+            if tasks:
+                await asyncio.gather(*tasks, return_exceptions=True)
+                logger.info("Web server tasks cancelled")
+                
                 if self.queue_manager and self.queue_manager.is_running:
                     await self.queue_manager.stop()
                     logger.info("Queue processor stopped")
 
-                # Stop training manager
+                if self.api_manager and self.api_manager.is_running:
+                    await self.api_manager.shutdown()
+                    logger.info("Stopped API health check loop")
+
                 if self.training_manager and self.training_manager.is_running:
                     await self.training_manager.stop()
                     logger.info("Training manager stopped")
@@ -378,19 +391,21 @@ def run_bot():
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
 
-    # Set up signal handler for SIGTERM only
+    # Set up signal handlers for both SIGTERM and SIGINT
     loop.add_signal_handler(
         signal.SIGTERM,
         lambda: asyncio.create_task(shutdown('SIGTERM', bot, loop))
     )
+    loop.add_signal_handler(
+        signal.SIGINT,
+        lambda: asyncio.create_task(shutdown('SIGINT', bot, loop))
+    )
     
-    # Restore default SIGINT handler
-    signal.signal(signal.SIGINT, signal.default_int_handler)
+    # Remove the line that sets the standard signal handler:
+    # signal.signal(signal.SIGINT, signal.default_int_handler)
 
     try:
         loop.run_until_complete(bot.start(token))
-    except KeyboardInterrupt:
-        logger.info("Received keyboard interrupt")
     except Exception as e:
         logger.error(f"Bot crashed: {e}")
         raise