diff options
| -rw-r--r-- | bot/__main__.py | 14 | ||||
| -rw-r--r-- | bot/utils/service_discovery.py | 22 |
2 files changed, 34 insertions, 2 deletions
diff --git a/bot/__main__.py b/bot/__main__.py index b1e9c61fa..d04be725c 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -7,6 +7,7 @@ from discord.ext.commands import Bot, when_mentioned_or from bot.constants import Bot as BotConfig, ClickUp from bot.formatter import Formatter +from bot.utils.service_discovery import wait_for_rmq log = logging.getLogger(__name__) @@ -36,13 +37,19 @@ bot.http_session = ClientSession( ) ) +log.info("Waiting for RabbitMQ...") +has_rmq = wait_for_rmq() + +if has_rmq: + log.info("RabbitMQ found") +else: + log.warning("Timed out while waiting for RabbitMQ") + # Internal/debug bot.load_extension("bot.cogs.logging") -bot.load_extension("bot.cogs.rmq") bot.load_extension("bot.cogs.security") bot.load_extension("bot.cogs.events") - # Commands, etc bot.load_extension("bot.cogs.bot") bot.load_extension("bot.cogs.cogs") @@ -64,6 +71,9 @@ bot.load_extension("bot.cogs.tags") bot.load_extension("bot.cogs.verification") bot.load_extension("bot.cogs.utils") +if has_rmq: + bot.load_extension("bot.cogs.rmq") + bot.run(BotConfig.token) bot.http_session.close() # Close the aiohttp session when the bot finishes running diff --git a/bot/utils/service_discovery.py b/bot/utils/service_discovery.py new file mode 100644 index 000000000..8d79096bd --- /dev/null +++ b/bot/utils/service_discovery.py @@ -0,0 +1,22 @@ +import datetime +import socket +import time +from contextlib import closing + +from bot.constants import RabbitMQ + +THIRTY_SECONDS = datetime.timedelta(seconds=30) + + +def wait_for_rmq(): + start = datetime.datetime.now() + + while True: + if datetime.datetime.now() - start > THIRTY_SECONDS: + return False + + with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock: + if sock.connect_ex((RabbitMQ.host, RabbitMQ.port)) == 0: + return True + + time.sleep(0.5) |