Enhance DiscordBot shutdown process; add connection timeout, cancel web tasks, and handle SIGINT signal
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user