slipstream/discordbot
slipstream/discordbot
1
1
Fork 0
Code Issues Pull Requests 3 Actions Packages Projects Releases Wiki Activity

feat: Enhance custom bot management with improved cleanup and resource handling

Browse Source
This commit is contained in:
Slipstream 2025-05-26 15:45:44 -06:00
parent f6e70a85c0
commit 2609c6ea8b
Signed by: slipstream
GPG Key ID: 13E498CE010AC6FD
3 changed files with 186 additions and 79 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

128
custom_bot_manager.py
View File

@ -45,6 +45,7 @@ class CustomBot(commands.Bot):
super().__init__(*args, **kwargs)
self.user_id = user_id
self.owner_id = int(os.getenv('OWNER_USER_ID', '0'))
self._cleanup_tasks = [] # Track cleanup tasks
async def setup_hook(self):
"""Called when the bot is first connected to Discord."""
@ -59,6 +60,26 @@ class CustomBot(commands.Bot):
log.error(f"Failed to load extension {cog} for custom bot {self.user_id}: {e}")
traceback.print_exc()
async def close(self):
"""Override close to ensure proper cleanup of all resources."""
log.info(f"Closing custom bot for user {self.user_id}...")
# Close all cogs that have aiohttp sessions
for cog_name, cog in self.cogs.items():
try:
if hasattr(cog, 'session') and cog.session and not cog.session.closed:
await cog.session.close()
log.info(f"Closed aiohttp session for cog {cog_name} in custom bot {self.user_id}")
except Exception as e:
log.error(f"Error closing session for cog {cog_name} in custom bot {self.user_id}: {e}")
# Wait a bit for sessions to close properly
await asyncio.sleep(0.1)
# Call parent close
await super().close()
log.info(f"Custom bot for user {self.user_id} closed successfully")
async def create_custom_bot(
user_id: str,
token: str,
@ -151,6 +172,13 @@ def run_custom_bot_in_thread(user_id: str, token: str) -> Tuple[bool, str]:
bot = custom_bots[user_id]
def _run_bot_thread():
"""Run the bot in a new event loop within this thread."""
try:
# Create a new event loop for this thread
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
async def _run_bot():
try:
await bot.start(token)
@ -162,11 +190,31 @@ def run_custom_bot_in_thread(user_id: str, token: str) -> Tuple[bool, str]:
log.error(f"Error running custom bot for user {user_id}: {e}")
custom_bot_status[user_id] = STATUS_ERROR
custom_bot_errors[user_id] = str(e)
finally:
# Ensure proper cleanup
if not bot.is_closed():
try:
await bot.close()
except Exception as e:
log.error(f"Error closing bot during cleanup for user {user_id}: {e}")
# Run the bot
loop.run_until_complete(_run_bot())
except Exception as e:
log.error(f"Error in bot thread for user {user_id}: {e}")
custom_bot_status[user_id] = STATUS_ERROR
custom_bot_errors[user_id] = str(e)
finally:
# Clean up the event loop
try:
loop.close()
except Exception as e:
log.error(f"Error closing event loop for user {user_id}: {e}")
# Create and start the thread
loop = asyncio.new_event_loop()
thread = threading.Thread(
target=lambda: loop.run_until_complete(_run_bot()),
target=_run_bot_thread,
daemon=True,
name=f"custom-bot-{user_id}"
)
@ -197,29 +245,51 @@ def stop_custom_bot(user_id: str) -> Tuple[bool, str]:
# Get the bot instance
bot = custom_bots[user_id]
# Close the bot (this will be done in a new thread to avoid blocking)
def _close_bot_thread():
"""Close the bot in a proper event loop context."""
try:
# Create a new event loop for this thread
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
async def _close_bot():
try:
await bot.close()
custom_bot_status[user_id] = STATUS_STOPPED
log.info(f"Successfully closed custom bot for user {user_id}")
except Exception as e:
log.error(f"Error closing custom bot for user {user_id}: {e}")
custom_bot_status[user_id] = STATUS_ERROR
custom_bot_errors[user_id] = str(e)
# Run the close operation
loop.run_until_complete(_close_bot())
except Exception as e:
log.error(f"Error in close thread for user {user_id}: {e}")
custom_bot_status[user_id] = STATUS_ERROR
custom_bot_errors[user_id] = str(e)
finally:
# Clean up the event loop
try:
loop.close()
except Exception as e:
log.error(f"Error closing event loop in close thread for user {user_id}: {e}")
# Run the close operation in a new thread
loop = asyncio.new_event_loop()
close_thread = threading.Thread(
target=lambda: loop.run_until_complete(_close_bot()),
target=_close_bot_thread,
daemon=True,
name=f"close-bot-{user_id}"
)
close_thread.start()
# Wait for the close thread to finish (with timeout)
close_thread.join(timeout=5.0)
close_thread.join(timeout=10.0) # Increased timeout for proper cleanup
# The thread will be cleaned up when the bot is started again
# Clean up the thread reference
if user_id in custom_bot_threads:
del custom_bot_threads[user_id]
return True, f"Stopped custom bot for user {user_id}"
@ -267,3 +337,47 @@ def get_all_custom_bot_statuses() -> Dict[str, Dict]:
for user_id in custom_bots:
result[user_id] = get_custom_bot_status(user_id)
return result
def list_custom_bots() -> List[Dict]:
"""
List all custom bot instances and their status.
Returns:
List of dictionaries containing bot information
"""
bots = []
for user_id in custom_bots.keys():
bot_info = get_custom_bot_status(user_id)
bots.append(bot_info)
return bots
def cleanup_all_custom_bots() -> None:
"""
Clean up all custom bot instances. Should be called when the main bot shuts down.
"""
log.info("Cleaning up all custom bots...")
# Stop all running bots
for user_id in list(custom_bots.keys()):
try:
if custom_bot_status.get(user_id) == STATUS_RUNNING:
log.info(f"Stopping custom bot for user {user_id}")
stop_custom_bot(user_id)
except Exception as e:
log.error(f"Error stopping custom bot for user {user_id}: {e}")
# Wait a bit for all bots to stop
import time
time.sleep(2)
# Force cleanup any remaining threads
for user_id, thread in list(custom_bot_threads.items()):
if thread.is_alive():
log.warning(f"Force terminating thread for custom bot {user_id}")
# Note: We can't force kill threads in Python, but we can clean up references
try:
del custom_bot_threads[user_id]
except KeyError:
pass
log.info("Custom bot cleanup completed")

8
discord_bot_sync_api.py
View File

@ -19,10 +19,18 @@ from typing import Optional # Added for GurtCog type hint
# --- Placeholder for GurtCog instance and bot instance ---
# These need to be set by the script that starts the bot and API server
# Import GurtCog and ModLogCog conditionally to avoid dependency issues
try:
from gurt.cog import GurtCog # Import GurtCog for type hint and access
from cogs.mod_log_cog import ModLogCog # Import ModLogCog for type hint
gurt_cog_instance: Optional[GurtCog] = None
mod_log_cog_instance: Optional[ModLogCog] = None # Placeholder for ModLogCog
except ImportError as e:
print(f"Warning: Could not import GurtCog or ModLogCog: {e}")
# Use Any type as fallback
from typing import Any
gurt_cog_instance: Optional[Any] = None
mod_log_cog_instance: Optional[Any] = None
bot_instance = None # Will be set to the Discord bot instance
# ============= Models =============

21
main.py
View File

@ -653,26 +653,11 @@ async def main(args): # Pass parsed args
else:
log.info("Flask server process was not running or already terminated.")
# Stop all custom bots
# Stop all custom bots using the improved cleanup function
try:
log.info("Stopping all custom bots...")
# Get all running custom bots
bot_statuses = custom_bot_manager.get_all_custom_bot_statuses()
stopped_count = 0
for user_id, status in bot_statuses.items():
if status.get('is_running', False):
log.info(f"Stopping custom bot for user {user_id}")
success, message = custom_bot_manager.stop_custom_bot(user_id)
if success:
stopped_count += 1
log.info(f"Successfully stopped custom bot for user {user_id}")
else:
log.error(f"Failed to stop custom bot for user {user_id}: {message}")
log.info(f"Stopped {stopped_count} custom bots")
custom_bot_manager.cleanup_all_custom_bots()
except Exception as e:
log.exception(f"Error stopping custom bots: {e}")
log.exception(f"Error during custom bot cleanup: {e}")
# Close database/cache pools if they were initialized
if bot.pg_pool: