Enhance DiscordBot shutdown process; add connection timeout, cancel web tasks, and handle SIGINT signal

This commit is contained in:
pacnpal
2025-02-25 13:09:09 -05:00
parent f40d19294f
commit 937bc18ca7

View File

@@ -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