RabbitMQ service discovery

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)